examples/node_classification/training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d089bbd1-1a16-4791-a28f-42bb6aefc540",
   "metadata": {},
   "source": [
    "# 1. Example of training GNNs for Node Classification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0ef1616-30c8-40ff-9ca3-67e11abc74e2",
   "metadata": {},
   "source": [
    "GRB provides easy-to-use APIs to train GNNs, facilitating the entire process from loading graph data, building GNN models, to evaluation and inference. Here is an example for the task of graph classification."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e6cec59-45da-46df-b840-28aa019d12e9",
   "metadata": {},
   "source": [
    "Contents\n",
    "- [Load Dataset](#Load-Dataset)\n",
    "- [Build Model](#Build-Model)\n",
    "- [Training](#Training)\n",
    "- [Inference](#Inference)\n",
    "- [Evaluation](#Evaluation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aa7e279c-5a47-427c-95b7-485d2f554fb7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:24.167412Z",
     "iopub.status.busy": "2021-09-26T13:00:24.167116Z",
     "iopub.status.idle": "2021-09-26T13:00:24.907032Z",
     "shell.execute_reply": "2021-09-26T13:00:24.906469Z",
     "shell.execute_reply.started": "2021-09-26T13:00:24.167338Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import grb.utils as utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e16e1d-262d-4e59-858c-b945af907603",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-11T11:59:56.178453Z",
     "iopub.status.busy": "2021-08-11T11:59:56.178207Z",
     "iopub.status.idle": "2021-08-11T11:59:56.181058Z",
     "shell.execute_reply": "2021-08-11T11:59:56.180559Z",
     "shell.execute_reply.started": "2021-08-11T11:59:56.178427Z"
    },
    "tags": []
   },
   "source": [
    "## 1.1. Load Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ba00b67-ef85-4e9d-a048-2f55f323f148",
   "metadata": {},
   "source": [
    "GRB datasets are named by the prefix *grb-*. There are four *mode* ('easy', 'medium', 'hard', 'full') for test set, representing different average degrees of test nodes, thus different difficulty for attacking them. The node features are processed by *arctan* normalization (first standardization then arctan function), which makes node features fall in the same scale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b6529982-a46b-4384-8b31-e61b93c0f0c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:25.398537Z",
     "iopub.status.busy": "2021-09-26T13:00:25.398307Z",
     "iopub.status.idle": "2021-09-26T13:00:25.440157Z",
     "shell.execute_reply": "2021-09-26T13:00:25.439614Z",
     "shell.execute_reply.started": "2021-09-26T13:00:25.398513Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset 'grb-cora' loaded.\n",
      "    Number of nodes: 2680\n",
      "    Number of edges: 5148\n",
      "    Number of features: 302\n",
      "    Number of classes: 7\n",
      "    Number of train samples: 1608\n",
      "    Number of val samples: 268\n",
      "    Number of test samples: 804\n",
      "    Dataset mode: full\n",
      "    Feature range: [-0.9406, 0.9430]\n"
     ]
    }
   ],
   "source": [
    "from grb.dataset import Dataset\n",
    "\n",
    "dataset_name = \"grb-cora\"\n",
    "dataset = Dataset(name=dataset_name, \n",
    "                  data_dir=\"../data/\",\n",
    "                  mode=\"full\",\n",
    "                  feat_norm=\"arctan\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97724892-f776-42ee-abcb-fadc554c5462",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1.2. Build Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df4fe302-e48d-4053-95aa-0ae65b8f46c6",
   "metadata": {},
   "source": [
    "GRB supports models based on pure Pytorch, CogDL or DGL. The following is an example of GCN implemented by pure Pytorch. Other models can be found in ``grb/model/torch``, ``grb/model/cogdl``, or ``grb/model/dgl``."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94a41133-70c9-481a-9e50-7d169a9a766b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.1. GCN (Graph Convolutional Network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7b490466-351a-43f8-b6c8-dfa1727b8f5b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:27.864610Z",
     "iopub.status.busy": "2021-09-26T13:00:27.864369Z",
     "iopub.status.idle": "2021-09-26T13:00:27.874485Z",
     "shell.execute_reply": "2021-09-26T13:00:27.874050Z",
     "shell.execute_reply.started": "2021-09-26T13:00:27.864582Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 24867.\n",
      "GCN(\n",
      "  (layers): ModuleList(\n",
      "    (0): LayerNorm((302,), eps=1e-05, elementwise_affine=True)\n",
      "    (1): GCNConv(\n",
      "      (linear): Linear(in_features=302, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
      "    (3): GCNConv(\n",
      "      (linear): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (4): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
      "    (5): GCNConv(\n",
      "      (linear): Linear(in_features=64, out_features=7, bias=True)\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import GCN\n",
    "from grb.utils.normalize import GCNAdjNorm\n",
    "\n",
    "model_name = \"gcn\"\n",
    "model = GCN(in_features=dataset.num_features,\n",
    "            out_features=dataset.num_classes,\n",
    "            hidden_features=64, \n",
    "            n_layers=3,\n",
    "            adj_norm_func=GCNAdjNorm,\n",
    "            layer_norm=True,\n",
    "            residual=False,\n",
    "            dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca7c3f7-1af3-4a6a-8546-3e18716427ef",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.2. GAT (Graph Attention Network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f1e7778e-b1f2-465c-93ef-b5ae49666f39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:29.479606Z",
     "iopub.status.busy": "2021-09-26T13:00:29.479362Z",
     "iopub.status.idle": "2021-09-26T13:00:29.626206Z",
     "shell.execute_reply": "2021-09-26T13:00:29.625635Z",
     "shell.execute_reply.started": "2021-09-26T13:00:29.479580Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 146197.\n",
      "GAT(\n",
      "  (layers): ModuleList(\n",
      "    (0): GATConv(\n",
      "      (fc): Linear(in_features=302, out_features=256, bias=False)\n",
      "      (feat_drop): Dropout(p=0.6, inplace=False)\n",
      "      (attn_drop): Dropout(p=0.6, inplace=False)\n",
      "      (leaky_relu): LeakyReLU(negative_slope=0.2)\n",
      "    )\n",
      "    (1): GATConv(\n",
      "      (fc): Linear(in_features=256, out_features=256, bias=False)\n",
      "      (feat_drop): Dropout(p=0.6, inplace=False)\n",
      "      (attn_drop): Dropout(p=0.6, inplace=False)\n",
      "      (leaky_relu): LeakyReLU(negative_slope=0.2)\n",
      "    )\n",
      "    (2): GATConv(\n",
      "      (fc): Linear(in_features=256, out_features=7, bias=False)\n",
      "      (feat_drop): Dropout(p=0.0, inplace=False)\n",
      "      (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "      (leaky_relu): LeakyReLU(negative_slope=0.2)\n",
      "    )\n",
      "  )\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ")\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using backend: pytorch\n"
     ]
    }
   ],
   "source": [
    "from grb.model.dgl import GAT\n",
    "\n",
    "model_name = \"gat\"\n",
    "model = GAT(in_features=dataset.num_features,\n",
    "            out_features=dataset.num_classes,\n",
    "            hidden_features=64,\n",
    "            n_layers=3,\n",
    "            n_heads=4,\n",
    "            adj_norm_func=None,\n",
    "            layer_norm=False,\n",
    "            residual=False,\n",
    "            feat_dropout=0.6,\n",
    "            attn_dropout=0.6,\n",
    "            dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7238600-e127-4ca7-b922-d6efc33b192c",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.3. APPNP (Approximated Personalized Propagation of Neural Predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "167c503b-6e4b-4004-894b-f4ba38ce7f18",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-17T15:58:47.938497Z",
     "iopub.status.busy": "2021-08-17T15:58:47.938403Z",
     "iopub.status.idle": "2021-08-17T15:58:47.942041Z",
     "shell.execute_reply": "2021-08-17T15:58:47.941753Z",
     "shell.execute_reply.started": "2021-08-17T15:58:47.938484Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 56199.\n",
      "APPNP(\n",
      "  (layers): ModuleList(\n",
      "    (0): Linear(in_features=302, out_features=128, bias=True)\n",
      "    (1): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (2): Linear(in_features=128, out_features=7, bias=True)\n",
      "  )\n",
      "  (edge_dropout): SparseEdgeDrop()\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import APPNP\n",
    "from grb.utils.normalize import GCNAdjNorm\n",
    "\n",
    "model_name = \"appnp\"\n",
    "model = APPNP(in_features=dataset.num_features,\n",
    "              out_features=dataset.num_classes,\n",
    "              hidden_features=128, \n",
    "              n_layers=3,\n",
    "              adj_norm_func=GCNAdjNorm,\n",
    "              layer_norm=False,\n",
    "              edge_drop=0.1,\n",
    "              alpha=0.01,\n",
    "              k=3,\n",
    "              dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "834e5f30-b800-4283-b00d-a0db823360da",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.4. GIN (Graph Isomorph Network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9a9cf67b-8702-4f7f-9b2e-8e9cb31ad272",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-17T15:58:47.942524Z",
     "iopub.status.busy": "2021-08-17T15:58:47.942439Z",
     "iopub.status.idle": "2021-08-17T15:58:47.950848Z",
     "shell.execute_reply": "2021-08-17T15:58:47.950446Z",
     "shell.execute_reply.started": "2021-08-17T15:58:47.942512Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 36745.\n",
      "GIN(\n",
      "  (layers): ModuleList(\n",
      "    (0): GINConv(\n",
      "      (linear1): Linear(in_features=302, out_features=64, bias=True)\n",
      "      (linear2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (norm): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (1): GINConv(\n",
      "      (linear1): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (linear2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (norm): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "  )\n",
      "  (mlp_layers): ModuleList(\n",
      "    (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (1): Linear(in_features=64, out_features=7, bias=True)\n",
      "  )\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import GIN\n",
    "\n",
    "model_name = \"gin\"\n",
    "model = GIN(in_features=dataset.num_features,\n",
    "            out_features=dataset.num_classes,\n",
    "            hidden_features=64, \n",
    "            n_layers=3,\n",
    "            adj_norm_func=None,\n",
    "            layer_norm=False,\n",
    "            batch_norm=True,\n",
    "            dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef4eb6ef-f241-42a7-a3e5-6d65ace2884a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.5. GraphSAGE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "35042c4a-27b4-46d1-aa14-4ee27df520ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-17T15:58:47.952387Z",
     "iopub.status.busy": "2021-08-17T15:58:47.952241Z",
     "iopub.status.idle": "2021-08-17T15:58:47.960127Z",
     "shell.execute_reply": "2021-08-17T15:58:47.959840Z",
     "shell.execute_reply.started": "2021-08-17T15:58:47.952368Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 147840.\n",
      "GraphSAGE(\n",
      "  (layers): ModuleList(\n",
      "    (0): SAGEConv(\n",
      "      (pool_layer): Linear(in_features=302, out_features=302, bias=True)\n",
      "      (linear1): Linear(in_features=302, out_features=64, bias=True)\n",
      "      (linear2): Linear(in_features=302, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (1): SAGEConv(\n",
      "      (pool_layer): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (linear1): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (linear2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (2): SAGEConv(\n",
      "      (pool_layer): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (linear1): Linear(in_features=64, out_features=7, bias=True)\n",
      "      (linear2): Linear(in_features=64, out_features=7, bias=True)\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import GraphSAGE\n",
    "from grb.utils.normalize import SAGEAdjNorm\n",
    "\n",
    "model_name = \"graphsage\"\n",
    "model = GraphSAGE(in_features=dataset.num_features,\n",
    "                  out_features=dataset.num_classes,\n",
    "                  hidden_features=64,\n",
    "                  n_layers=3,\n",
    "                  adj_norm_func=SAGEAdjNorm,\n",
    "                  layer_norm=False,\n",
    "                  dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7464f5b2-02d8-4ffd-b688-7e450930f408",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.6. SGCN (Simplified Graph Convolutional Network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fb3ae0dc-94c5-4bf3-a745-a00353d9dbf6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-17T15:58:47.960705Z",
     "iopub.status.busy": "2021-08-17T15:58:47.960621Z",
     "iopub.status.idle": "2021-08-17T15:58:47.963959Z",
     "shell.execute_reply": "2021-08-17T15:58:47.963611Z",
     "shell.execute_reply.started": "2021-08-17T15:58:47.960694Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 24611.\n",
      "SGCN(\n",
      "  (batch_norm): BatchNorm1d(302, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (in_conv): Linear(in_features=302, out_features=64, bias=True)\n",
      "  (out_conv): Linear(in_features=64, out_features=7, bias=True)\n",
      "  (layers): ModuleList(\n",
      "    (0): SGConv(\n",
      "      (linear): Linear(in_features=64, out_features=64, bias=True)\n",
      "    )\n",
      "  )\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import SGCN\n",
    "from grb.utils.normalize import GCNAdjNorm\n",
    "\n",
    "model_name = \"sgcn\"\n",
    "model = SGCN(in_features=dataset.num_features,\n",
    "             out_features=dataset.num_classes,\n",
    "             hidden_features=64,\n",
    "             n_layers=3,\n",
    "             adj_norm_func=GCNAdjNorm,\n",
    "             k=4,\n",
    "             dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c00964fa-a0e3-4995-8bb2-7c8c8c165e01",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.7. TAGCN (Topological Adaptive Graph Convolutional Network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "12b57a62-a650-42fa-84a7-57d1cfa066ef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-17T15:58:47.964476Z",
     "iopub.status.busy": "2021-08-17T15:58:47.964394Z",
     "iopub.status.idle": "2021-08-17T15:58:47.968369Z",
     "shell.execute_reply": "2021-08-17T15:58:47.968044Z",
     "shell.execute_reply.started": "2021-08-17T15:58:47.964465Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 71751.\n",
      "TAGCN(\n",
      "  (layers): ModuleList(\n",
      "    (0): TAGConv(\n",
      "      (linear): Linear(in_features=906, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (1): TAGConv(\n",
      "      (linear): Linear(in_features=192, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (2): TAGConv(\n",
      "      (linear): Linear(in_features=192, out_features=7, bias=True)\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import TAGCN\n",
    "from grb.utils.normalize import GCNAdjNorm\n",
    "\n",
    "model_name = \"tagcn\"\n",
    "model = TAGCN(in_features=dataset.num_features,\n",
    "              out_features=dataset.num_classes,\n",
    "              hidden_features=64,\n",
    "              n_layers=3,\n",
    "              adj_norm_func=GCNAdjNorm,\n",
    "              k=2,\n",
    "              dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eef66d7-2e06-4d15-a557-58d232d116e7",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.2.8. MLP (Multi-layer Perceptron)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d6d0662a-7c52-48d3-a714-fcbe7ba17452",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-08-17T15:58:47.969099Z",
     "iopub.status.busy": "2021-08-17T15:58:47.968960Z",
     "iopub.status.idle": "2021-08-17T15:58:47.973513Z",
     "shell.execute_reply": "2021-08-17T15:58:47.973136Z",
     "shell.execute_reply.started": "2021-08-17T15:58:47.969081Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 24007.\n",
      "MLP(\n",
      "  (layers): ModuleList(\n",
      "    (0): MLPLayer(\n",
      "      (linear): Linear(in_features=302, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (1): MLPLayer(\n",
      "      (linear): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (2): MLPLayer(\n",
      "      (linear): Linear(in_features=64, out_features=7, bias=True)\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from grb.model.torch import MLP\n",
    "\n",
    "model_name = \"mlp\"\n",
    "model = MLP(in_features=dataset.num_features,\n",
    "            out_features=dataset.num_classes,\n",
    "            hidden_features=64, \n",
    "            n_layers=3,\n",
    "            dropout=0.5)\n",
    "print(\"Number of parameters: {}.\".format(utils.get_num_params(model)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee0a75f-8620-4c42-abc8-72d9bc65d18c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1.3. Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09e37b87-01ee-43da-9dab-52f7b34f842c",
   "metadata": {},
   "source": [
    "GRB provides ``grb.utils.trainer`` that facilitates the training process of GNNs. The training mode can be chosen from ``inductive`` or ``transductive``. In the inductive mode, only train nodes can be seen during training, train+val nodes can be seen during validation, train+val+test nodes can be seen during testing. In the transductive mode, all nodes are available for each process.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "27e138eb-8190-4f3c-affa-c55053e22cb0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:41.808713Z",
     "iopub.status.busy": "2021-09-26T13:00:41.808469Z",
     "iopub.status.idle": "2021-09-26T13:00:41.812170Z",
     "shell.execute_reply": "2021-09-26T13:00:41.811582Z",
     "shell.execute_reply.started": "2021-09-26T13:00:41.808687Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_dir = \"./saved_models/{}/{}\".format(dataset_name, model_name)\n",
    "save_name = \"model.pt\"\n",
    "device = \"cuda:0\"\n",
    "feat_norm = None\n",
    "train_mode = \"inductive\"  # \"transductive\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "43c9c619-cd60-47a3-8d9d-55d6c8525c32",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:42.431621Z",
     "iopub.status.busy": "2021-09-26T13:00:42.431390Z",
     "iopub.status.idle": "2021-09-26T13:00:44.584956Z",
     "shell.execute_reply": "2021-09-26T13:00:44.584323Z",
     "shell.execute_reply.started": "2021-09-26T13:00:42.431596Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from grb.trainer.trainer import Trainer\n",
    "\n",
    "trainer = Trainer(dataset=dataset, \n",
    "                  optimizer=torch.optim.Adam(model.parameters(), lr=0.01),\n",
    "                  loss=torch.nn.functional.cross_entropy,\n",
    "                  lr_scheduler=False,\n",
    "                  early_stop=True,\n",
    "                  early_stop_patience=500,\n",
    "                  feat_norm=feat_norm,\n",
    "                  device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f567b85c-9a14-46e1-9a8f-17c94c0080bb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:00:45.463271Z",
     "iopub.status.busy": "2021-09-26T13:00:45.463031Z",
     "iopub.status.idle": "2021-09-26T13:01:05.425966Z",
     "shell.execute_reply": "2021-09-26T13:01:05.425642Z",
     "shell.execute_reply.started": "2021-09-26T13:00:45.463245Z"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef9d1a0fe6184004a547659f40613e3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training early stopped. Best validation score: 0.8134\n",
      "Training runtime: 19.9552.\n"
     ]
    }
   ],
   "source": [
    "trainer.train(model=model, \n",
    "              n_epoch=2000,\n",
    "              eval_every=1,\n",
    "              save_after=0,\n",
    "              save_dir=save_dir,\n",
    "              save_name=save_name,\n",
    "              train_mode=train_mode,\n",
    "              verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23423752-5dd9-4f8a-89b8-8de74e21b92c",
   "metadata": {},
   "source": [
    "## 1.4. Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4d0cf8b1-4518-46bf-b92f-378bb976641a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:01:09.045438Z",
     "iopub.status.busy": "2021-09-26T13:01:09.045185Z",
     "iopub.status.idle": "2021-09-26T13:01:09.055946Z",
     "shell.execute_reply": "2021-09-26T13:01:09.055323Z",
     "shell.execute_reply.started": "2021-09-26T13:01:09.045411Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GAT(\n",
       "  (layers): ModuleList(\n",
       "    (0): GATConv(\n",
       "      (fc): Linear(in_features=302, out_features=256, bias=False)\n",
       "      (feat_drop): Dropout(p=0.6, inplace=False)\n",
       "      (attn_drop): Dropout(p=0.6, inplace=False)\n",
       "      (leaky_relu): LeakyReLU(negative_slope=0.2)\n",
       "    )\n",
       "    (1): GATConv(\n",
       "      (fc): Linear(in_features=256, out_features=256, bias=False)\n",
       "      (feat_drop): Dropout(p=0.6, inplace=False)\n",
       "      (attn_drop): Dropout(p=0.6, inplace=False)\n",
       "      (leaky_relu): LeakyReLU(negative_slope=0.2)\n",
       "    )\n",
       "    (2): GATConv(\n",
       "      (fc): Linear(in_features=256, out_features=7, bias=False)\n",
       "      (feat_drop): Dropout(p=0.0, inplace=False)\n",
       "      (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "      (leaky_relu): LeakyReLU(negative_slope=0.2)\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.5, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = torch.load(os.path.join(save_dir, save_name))\n",
    "model = model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0770136c-3259-4e51-b4ae-e2fc28a5e763",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:01:09.698847Z",
     "iopub.status.busy": "2021-09-26T13:01:09.698638Z",
     "iopub.status.idle": "2021-09-26T13:01:09.709520Z",
     "shell.execute_reply": "2021-09-26T13:01:09.709019Z",
     "shell.execute_reply.started": "2021-09-26T13:01:09.698826Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# by trainer\n",
    "pred = trainer.inference(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "278be10b-440a-4068-be2e-237e8d0b2e7e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:01:10.430075Z",
     "iopub.status.busy": "2021-09-26T13:01:10.429895Z",
     "iopub.status.idle": "2021-09-26T13:01:10.441568Z",
     "shell.execute_reply": "2021-09-26T13:01:10.441095Z",
     "shell.execute_reply.started": "2021-09-26T13:01:10.430053Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# by utils\n",
    "pred = utils.inference(model, \n",
    "                       features=dataset.features,\n",
    "                       feat_norm=feat_norm,\n",
    "                       adj=dataset.adj,\n",
    "                       adj_norm_func=model.adj_norm_func,\n",
    "                       device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d1df58f-7116-47c0-8e1d-c70b0567c964",
   "metadata": {},
   "source": [
    "## 1.5. Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bd7d828e-029d-4127-a6ee-37cdf316c59a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:01:11.527354Z",
     "iopub.status.busy": "2021-09-26T13:01:11.527097Z",
     "iopub.status.idle": "2021-09-26T13:01:11.538947Z",
     "shell.execute_reply": "2021-09-26T13:01:11.538526Z",
     "shell.execute_reply.started": "2021-09-26T13:01:11.527329Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test score: 0.8458\n"
     ]
    }
   ],
   "source": [
    "# by trainer\n",
    "test_score = trainer.evaluate(model, dataset.test_mask)\n",
    "print(\"Test score: {:.4f}\".format(test_score))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e3cf4e0a-afae-464a-9893-d2760b4734b4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-09-26T13:01:12.159340Z",
     "iopub.status.busy": "2021-09-26T13:01:12.159083Z",
     "iopub.status.idle": "2021-09-26T13:01:12.172066Z",
     "shell.execute_reply": "2021-09-26T13:01:12.171555Z",
     "shell.execute_reply.started": "2021-09-26T13:01:12.159313Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test score: 0.8458\n"
     ]
    }
   ],
   "source": [
    "# by utils\n",
    "test_score = utils.evaluate(model, \n",
    "                            features=dataset.features,\n",
    "                            adj=dataset.adj,\n",
    "                            labels=dataset.labels,\n",
    "                            feat_norm=feat_norm,\n",
    "                            adj_norm_func=model.adj_norm_func,\n",
    "                            mask=dataset.test_mask,\n",
    "                            device=device)\n",
    "print(\"Test score: {:.4f}\".format(test_score))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25c1d781-8ad3-4f64-b59c-ce72c68d30d2",
   "metadata": {},
   "source": [
    "For further information, please refer to the [GRB Documentation](https://grb.readthedocs.io/en/latest/)."
   ]
  }
 ],
 "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.11"
  },
  "toc-autonumbering": true,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}