gcn_baseline_wip.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch_geometric\n",
    "import networkx as nx\n",
    "import os\n",
    "import torch\n",
    "from torch_geometric.loader import DataLoader\n",
    "from tqdm import tqdm\n",
    "from torch_geometric.data import Data\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#>>> pd.read_excel('data/RhythmNames.xlsx')\n",
    "#   Acronym Name                                     Full Name\n",
    "#0          SB                               Sinus Bradycardia\n",
    "#1          SR                                    Sinus Rhythm\n",
    "#2         AFIB                            Atrial Fibrillation\n",
    "#3         ST                                Sinus Tachycardia\n",
    "#4          AF                                  Atrial Flutter\n",
    "#5          SI                              Sinus Irregularity\n",
    "#6          SVT                   Supraventricular Tachycardia\n",
    "#7          AT                              Atrial Tachycardia\n",
    "#8         AVNRT  Atrioventricular  Node Reentrant Tachycardia\n",
    "#9          AVRT        Atrioventricular Reentrant Tachycardia\n",
    "#10        SAAWR       Sinus Atrium to Atrial Wandering Rhythm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#>>> pd.read_excel('data/ConditionNames.xlsx')\n",
    "#   Acronym Name                                  Full Name\n",
    "#0          1AVB            1 degree atrioventricular block\n",
    "#1          2AVB            2 degree atrioventricular block\n",
    "#2         2AVB1  2 degree atrioventricular block(Type one)\n",
    "#3         2AVB2  2 degree atrioventricular block(Type two)\n",
    "#4          3AVB            3 degree atrioventricular block\n",
    "#5           ABI                            atrial bigeminy\n",
    "#6           ALS                            Axis left shift\n",
    "#7           APB                     atrial premature beats\n",
    "#8           AQW                            abnormal Q wave\n",
    "#9           ARS                           Axis right shift\n",
    "#10          AVB                     atrioventricular block\n",
    "#11          CCR                 countercolockwise rotation\n",
    "#12           CR                        colockwise rotation\n",
    "#13          ERV     Early repolarization of the ventricles\n",
    "#14         FQRS                                  fQRS Wave\n",
    "#15          IDC            Interior differences conduction\n",
    "#16          IVB                     Intraventricular block\n",
    "#17          JEB                     junctional escape beat\n",
    "#18          JPS                              J point shift\n",
    "#19          JPT                  junctional premature beat\n",
    "#20         LBBB                   left bundle branch block\n",
    "#21        LBBBB              left back bundle branch block\n",
    "#22        LFBBB             left front bundle branch block\n",
    "#23         LRRI                           Long RR interval\n",
    "#24          LVH                 left ventricle hypertrophy\n",
    "#25         LVHV                left ventricle high voltage\n",
    "#26      LVQRSAL              lower voltage QRS in all lead\n",
    "#27      LVQRSCL            lower voltage QRS in chest lead\n",
    "#28      LVQRSLL             lower voltage QRS in limb lead\n",
    "#29           MI                      myocardial infarction\n",
    "#30         MIBW         myocardial infraction in back wall\n",
    "#31         MIFW   Myocardial infgraction in the front wall\n",
    "#32         MILW    Myocardial infraction in the lower wall\n",
    "#33         MISW     Myocardial infraction in the side wall\n",
    "#34         PRIE                      PR interval extension\n",
    "#35          PWC                              P wave Change\n",
    "#36         QTIE                      QT interval extension\n",
    "#37          RAH                   right atrial hypertrophy\n",
    "#38         RAHV                  right atrial high voltage\n",
    "#39         RBBB                  right bundle branch block\n",
    "#40          RVH                right ventricle hypertrophy\n",
    "#41         STDD                               ST drop down\n",
    "#42          STE                               ST extension\n",
    "#43         STTC                                ST-T Change\n",
    "#44         STTU                                 ST tilt up\n",
    "#45          TWC                              T wave Change\n",
    "#46          TWO                            T wave opposite\n",
    "#47           UW                                     U wave\n",
    "#48           VB                       ventricular bigeminy\n",
    "#49          VEB                    ventricular escape beat\n",
    "#50          VFW                    ventricular fusion wave\n",
    "#51          VPB                 ventricular premature beat\n",
    "#52          VPE                  ventricular preexcitation\n",
    "#53          VET               ventricular escape trigeminy\n",
    "#54         WAVN   Wandering in the atrioventricalualr node\n",
    "#55          WPW                                        WPW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_edge_index(node_num=12):\n",
    "    src_nodes = [idx for idx in range(node_num) for _ in range(node_num)]\n",
    "    tgt_nodes = [idx for _ in range(node_num) for idx in range(node_num)]\n",
    "    edge_index = torch.tensor(np.array([src_nodes, tgt_nodes]), dtype=torch.long)\n",
    "    return edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagnostics_df = pd.read_excel('data/Diagnostics.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import LabelEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LabelEncoder()"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "le = LabelEncoder()\n",
    "le.fit(diagnostics_df.Rhythm.unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['AF', 'AFIB', 'AT', 'AVNRT', 'AVRT', 'SA', 'SAAWR', 'SB', 'SR',\n",
       "       'ST', 'SVT'], dtype=object)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "le.classes_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagnostics_df['label'] = le.transform(diagnostics_df.Rhythm.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0         1\n",
       "1         7\n",
       "2         5\n",
       "3         7\n",
       "4         0\n",
       "         ..\n",
       "10641    10\n",
       "10642    10\n",
       "10643    10\n",
       "10644    10\n",
       "10645    10\n",
       "Name: label, Length: 10646, dtype: int64"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diagnostics_df['label']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 5675/10653 [03:20<04:37, 17.96it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MUSE_20180113_124215_52000.csv\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10653/10653 [06:21<00:00, 27.95it/s]\n"
     ]
    }
   ],
   "source": [
    "data_list = []\n",
    "for file_name in tqdm(os.listdir('data/')):\n",
    "    if file_name.endswith('.csv'):\n",
    "        sub_name = file_name.split('.')[0]\n",
    "        y = torch.tensor(diagnostics_df.loc[diagnostics_df.FileName == sub_name].label.values[0])\n",
    "        df = pd.read_csv(f'data/{file_name}', header=None)\n",
    "        if len(df) != 5000:\n",
    "            print(file_name)\n",
    "            continue\n",
    "        x = torch.Tensor(np.array([df[col].values for col in df.columns]))\n",
    "        edge_index = get_edge_index(node_num=12)\n",
    "        data = Data(x=x, edge_index=edge_index, y=y)\n",
    "        data_list.append(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(data_list[:8000], batch_size=32, shuffle=True)\n",
    "test_loader = DataLoader(data_list[8000:], batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance = next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "#g = torch_geometric.utils.to_networkx(instance, to_undirected=True)\n",
    "#nx.draw(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(8)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instance.y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 1:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 2:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 3:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 4:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 5:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 6:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 7:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 8:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 9:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 10:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 11:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 12:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 13:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 14:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 15:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 16:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 17:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 18:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 19:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 20:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 21:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 22:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 23:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 24:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 25:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 26:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 27:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 28:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 29:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 30:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 31:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 32:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 33:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 34:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 35:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 36:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 37:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 38:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 39:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 40:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 41:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 42:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 43:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 44:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 45:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 46:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 47:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 48:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 49:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 50:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 51:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 52:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 53:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 54:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 55:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 56:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 57:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 58:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 59:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 60:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 61:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 62:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 63:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 64:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 65:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 66:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 67:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 68:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 69:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 70:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 71:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 72:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 73:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 74:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 75:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 76:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 77:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 78:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 79:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 80:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 81:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 82:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 83:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 84:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 85:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 86:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 87:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 88:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 89:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 90:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 91:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 92:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 93:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 94:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 95:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 96:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 97:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 98:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 99:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 100:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 101:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 102:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 103:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 104:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 105:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 106:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 107:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 108:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 109:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 110:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 111:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 112:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 113:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 114:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 115:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 116:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 117:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 118:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 119:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 120:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 121:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 122:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 123:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 124:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 125:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 126:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 127:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 128:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 129:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 130:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 131:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 132:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 133:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 134:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 135:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 136:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 137:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 138:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 139:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 140:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 141:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 142:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 143:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 144:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 145:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 146:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 147:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 148:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 149:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 150:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 151:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 152:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 153:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 154:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 155:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 156:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 157:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 158:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 159:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 160:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 161:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 162:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 163:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 164:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 165:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 166:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 167:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 168:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 169:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 170:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 171:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 172:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 173:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 174:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 175:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 176:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 177:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 178:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 179:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 180:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 181:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 182:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 183:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 184:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 185:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 186:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 187:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 188:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 189:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 190:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 191:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 192:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 193:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 194:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 195:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 196:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 197:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 198:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 199:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 200:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 201:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 202:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 203:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 204:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 205:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 206:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 207:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 208:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 209:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 210:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 211:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 212:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 213:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 214:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 215:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 216:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 217:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 218:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 219:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 220:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 221:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 222:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 223:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 224:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 225:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 226:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 227:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 228:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 229:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 230:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 231:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 232:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 233:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 234:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 235:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 236:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 237:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 238:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 239:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 240:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 241:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 242:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 243:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 244:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 245:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 246:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 247:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 248:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 249:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n",
      "Step 250:\n",
      "=======\n",
      "Number of graphs in the current batch: 32\n",
      "DataBatch(x=[384, 5000], edge_index=[2, 4608], y=[32], batch=[384], ptr=[33])\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for step, data in enumerate(train_loader):\n",
    "    print(f'Step {step + 1}:')\n",
    "    print('=======')\n",
    "    print(f'Number of graphs in the current batch: {data.num_graphs}')\n",
    "    print(data)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_channels = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GCN(\n",
      "  (conv1): GCNConv(-1, 250)\n",
      "  (conv2): GCNConv(250, 250)\n",
      "  (conv3): GCNConv(250, 250)\n",
      "  (lin): Linear(in_features=250, out_features=11, bias=True)\n",
      ")\n",
      "GAT(\n",
      "  (conv1): GATConv(-1, 250, heads=1)\n",
      "  (conv2): GATConv(250, 250, heads=1)\n",
      "  (conv3): GATConv(250, 250, heads=1)\n",
      "  (lin): Linear(in_features=250, out_features=11, bias=True)\n",
      ")\n",
      "GNN(\n",
      "  (conv1): GraphConv(-1, 250)\n",
      "  (conv2): GraphConv(250, 250)\n",
      "  (conv3): GraphConv(250, 250)\n",
      "  (lin): Linear(in_features=250, out_features=11, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Training Graph Neural Networks for Graph Classification\n",
    "# Embed each node by performing multiple rounds of message passing\n",
    "# Aggregate node embeddings into a unified graph embedding (readout layer)\n",
    "# Train a final classifier on the graph embedding\n",
    " \n",
    "# There exists multiple readout layers in literature, but the most common one is to simply take the average of node embeddings\n",
    "\n",
    "# For ex: In GCNConv we use the  ReLU(𝑥)=max(𝑥,0)  activation for obtaining localized node embeddings,\n",
    "# before we apply our final classifier on top of a graph readout layer.\n",
    "\n",
    "# PyTorch Geometric provides this functionality via torch_geometric.nn.global_mean_pool, \n",
    "# which takes in the node embeddings of all nodes in the mini-batch and the assignment vector batch \n",
    "# to compute a graph embedding of size [batch_size, hidden_channels] for each graph in the batch.\n",
    "\n",
    "# The final architecture for applying GNNs to the task of graph classification then looks as follows and allows for complete end-to-end training:\n",
    "\n",
    "from torch.nn import Linear\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv, GATConv\n",
    "from torch_geometric.nn import global_mean_pool\n",
    "from torch_geometric.nn import GraphConv\n",
    "\n",
    "\n",
    "class GNN(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super(GNN, self).__init__()\n",
    "        torch.manual_seed(12345)\n",
    "        self.conv1 = GraphConv(-1, hidden_channels)\n",
    "        self.conv2 = GraphConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GraphConv(hidden_channels, hidden_channels)\n",
    "        self.lin = Linear(hidden_channels, len(le.classes_))\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index)\n",
    "\n",
    "        x = global_mean_pool(x, batch)\n",
    "\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super(GCN, self).__init__()\n",
    "        torch.manual_seed(12345)\n",
    "        self.conv1 = GCNConv(-1, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.lin = Linear(hidden_channels, len(le.classes_))\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        # 1. Obtain node embeddings \n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index)\n",
    "\n",
    "        # 2. Readout layer\n",
    "        x = global_mean_pool(x, batch)  # [batch_size, hidden_channels]\n",
    "\n",
    "        # 3. Apply a final classifier\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "\n",
    "class GAT(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super(GAT, self).__init__()\n",
    "        torch.manual_seed(12345)\n",
    "        self.conv1 = GATConv(-1, hidden_channels)\n",
    "        self.conv2 = GATConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GATConv(hidden_channels, hidden_channels)\n",
    "        self.lin = Linear(hidden_channels, len(le.classes_))\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        # 1. Obtain node embeddings \n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index)\n",
    "\n",
    "        # 2. Readout layer\n",
    "        x = global_mean_pool(x, batch)  # [batch_size, hidden_channels]\n",
    "\n",
    "        # 3. Apply a final classifier\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin(x)\n",
    "        \n",
    "        return x\n",
    "       \n",
    "# Define the models\n",
    "model1 = GCN(hidden_channels=hidden_channels)\n",
    "model2= GAT(hidden_channels=hidden_channels)\n",
    "model3 = GNN(hidden_channels=hidden_channels)\n",
    "\n",
    "# Print them \n",
    "print(model1)\n",
    "print(model2)\n",
    "print(model3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set model paramters and model type\n",
    "def set_model_parameters(model_type, lr=0.01):\n",
    "    model = model_type\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    return model, optimizer, criterion\n",
    "\n",
    "# Train the model\n",
    "def train(model, optimizer,criterion):\n",
    "    model.train()\n",
    "\n",
    "    for data in train_loader:  # Iterate in batches over the training dataset.\n",
    "         out = model(data.x, data.edge_index, data.batch)  # Perform a single forward pass.\n",
    "         loss = criterion(out, data.y)  # Compute the loss.\n",
    "         loss.backward()  # Derive gradients.\n",
    "         optimizer.step()  # Update parameters based on gradients.\n",
    "         optimizer.zero_grad()  # Clear gradients.\n",
    "\n",
    "# Test the model \n",
    "def test(loader, model):\n",
    "     model.eval()\n",
    "\n",
    "     correct = 0\n",
    "     for data in loader:  # Iterate in batches over the training/test dataset.\n",
    "         out = model(data.x, data.edge_index, data.batch)  \n",
    "         pred = out.argmax(dim=1)  # Use the class with highest probability.\n",
    "         correct += int((pred == data.y).sum())  # Check against ground-truth labels.\n",
    "     return correct / len(loader.dataset)  # Derive ratio of correct predictions.\n",
    "\n",
    "# Training and Testing Pipeline \n",
    "def running_epochs(model,optimizer,criterion):\n",
    "    for epoch in range(1, 10):\n",
    "        train(model,optimizer,criterion)\n",
    "        train_acc = test(train_loader, model)\n",
    "        test_acc = test(test_loader, model)\n",
    "        print(f'Epoch: {epoch:03d}, Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Acc: 0.0428, Test Acc: 0.0389\n",
      "Epoch: 002, Train Acc: 0.0428, Test Acc: 0.0389\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-40-cd2968c70205>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Experiment 1: GCN Baseline\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mset_model_parameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.001\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mrunning_epochs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mcriterion\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-38-e55c55e75cb2>\u001b[0m in \u001b[0;36mrunning_epochs\u001b[0;34m(model, optimizer, criterion)\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mrunning_epochs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mcriterion\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 33\u001b[0;31m         \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mcriterion\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     34\u001b[0m         \u001b[0mtrain_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m         \u001b[0mtest_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-38-e55c55e75cb2>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(model, optimizer, criterion)\u001b[0m\n\u001b[1;32m     13\u001b[0m          \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Perform a single forward pass.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m          \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Compute the loss.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m          \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Derive gradients.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m          \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Update parameters based on gradients.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m          \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Clear gradients.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    361\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    362\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 363\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    364\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    365\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    171\u001b[0m     \u001b[0;31m# some Python versions print out the first line of a multi-line function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    172\u001b[0m     \u001b[0;31m# calls in the traceback and some print out the last line\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 173\u001b[0;31m     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[1;32m    174\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    175\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Experiment 1: GCN Baseline \n",
    "model, optimizer, criterion = set_model_parameters(model1, lr=0.001)\n",
    "running_epochs(model,optimizer,criterion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}