GNNs/PyG/gcn_node_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Convolutional Network\n",
    "This notebook demonstrates the training of [Graph Convolutional Networks (GCN)](https://arxiv.org/pdf/1609.02907.pdf) with TigerGraph. [Pytorch Geometric](https://pytorch-geometric.readthedocs.io)'s implementation of GCN is used here. We train the model on the Cora dataset from [PyG datasets](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.Planetoid) with TigerGraph as the data store. The dataset contains 2708 machine learning papers and 10556 citation links between the papers.  Each publication in the dataset is described by a 0/1-valued word vector indicating the absence/presence of the corresponding word from a dictionary. The dictionary consists of 1433 unique words. Each paper is classified into one of seven classes based on the topic. The goal is to predict the class of each vertex in the graph.\n",
    "\n",
    "The following libraries are required to run this notebook. Uncomment to install them if necessary. You might need to restart the kernel after installing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install torch==1.12.0 --extra-index-url https://download.pytorch.org/whl/cpu\n",
    "#!pip install torch-scatter==2.0.9 torch-sparse==0.6.14 torch-cluster==1.6.0 torch-spline-conv==1.2.1 torch-geometric==2.0.4 -f https://data.pyg.org/whl/torch-1.12.0+cpu.html\n",
    "#!pip install pyTigerGraph[gds]\n",
    "#!pip install tensorboard # If you use tensorboard for visualization later"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "* [Data Processing](#data_processing)  \n",
    "* [Train on whole graph](#train_whole)  \n",
    "* [Train on neighborhood subgraphs](#train_subgraph)  \n",
    "* [Inference](#inference)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Processing <a name=\"data_processing\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connect to TigerGraph\n",
    "\n",
    "The `TigerGraphConnection` class represents a connection to the TigerGraph database. Under the hood, it stores the necessary information to communicate with the database. It is able to perform quite a few database tasks. Please see its [documentation](https://docs.tigergraph.com/pytigergraph/current/intro/) for details. \n",
    "\n",
    "To connect your database, modify the `config.json` file accompanying this notebook. Set the value of `getToken` based on whether token auth is enabled for your database. Token auth is always enabled for tgcloud databases. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyTigerGraph import TigerGraphConnection\n",
    "import json\n",
    "\n",
    "# Read in DB configs\n",
    "with open('../../config.json', \"r\") as config_file:\n",
    "    config = json.load(config_file)\n",
    "    \n",
    "conn = TigerGraphConnection(\n",
    "    host=config[\"host\"],\n",
    "    username=config[\"username\"],\n",
    "    password=config[\"password\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ingest Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca5fe2d80a154927ab6a6cc5df719e00",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0/166537 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from pyTigerGraph.datasets import Datasets\n",
    "\n",
    "dataset = Datasets(\"Cora\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Checking database ----\n",
      "A graph with name Cora already exists in the database. Please drop it first before ingesting.\n"
     ]
    }
   ],
   "source": [
    "conn.ingestDataset(dataset, getToken=config[\"getToken\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize Schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94d8603d9be242909ba1ddaef2c45cde",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "CytoscapeWidget(cytoscape_layout={'name': 'circle', 'animate': True, 'padding': 1}, cytoscape_style=[{'selecto…"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pyTigerGraph.visualization import drawSchema\n",
    "\n",
    "drawSchema(conn.getSchema(force=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Paper': 2708}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conn.getVertexCount('*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Cite': 10556}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conn.getEdgeCount()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train/validation/test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The code in this cell is commented out because there is no need to split the vertices into \n",
    "# training/validation/test sets, as the split is already done in the original dataset. \n",
    "# See notebooks under basics for examples on the split function.\n",
    "\n",
    "#split = conn.gds.vertexSplitter(train_mask=0.8, val_mask=0.1, test_mask=0.1)\n",
    "#split.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of vertices in training set: 140\n",
      "Number of vertices in validation set: 500\n",
      "Number of vertices in test set: 1000\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Number of vertices in training set:\",\n",
    "    conn.getVertexCount(\"Paper\", where=\"train_mask!=0\"),\n",
    ")\n",
    "print(\n",
    "    \"Number of vertices in validation set:\",\n",
    "    conn.getVertexCount(\"Paper\", where=\"val_mask!=0\"),\n",
    ")\n",
    "print(\n",
    "    \"Number of vertices in test set:\", \n",
    "    conn.getVertexCount(\"Paper\", where=\"test_mask!=0\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on whole graph <a name=\"train_whole\"></a>\n",
    "We first train the model on the whole graph. This will **NOT** work when the graph is large. See the section of training on subgraphs for real use. However, we still include this example for illustration purpose. Hyperparameters for the model and training environment are defined below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "hp = {\"hidden_dim\": 64, \"num_layers\": 2,\n",
    "      \"dropout\": 0.6, \"lr\": 0.01, \"l2_penalty\": 5e-4}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct graph loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `GraphLoader` can get the whole graph from database all at once (`num_batches=1`). See the tutorial on dataloaders under basics for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing and optimizing queries. It might take a minute if this is the first time you use this loader.\n",
      "Query installation finished.\n"
     ]
    }
   ],
   "source": [
    "graph_loader = conn.gds.graphLoader(\n",
    "    v_in_feats=[\"x\"],\n",
    "    v_out_labels=[\"y\"],\n",
    "    v_extra_feats=[\"train_mask\", \"val_mask\", \"test_mask\"],\n",
    "    num_batches=1,\n",
    "    output_format=\"PyG\",\n",
    "    shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 10556], x=[2708, 1433], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the whole graph from the loader\n",
    "data = graph_loader.data\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct model and optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We build a GCN model with 2 convolutional layers, and use the Adam optimizer with a learning rate of 0.01."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = GCN(\n",
    "    in_channels=1433,\n",
    "    hidden_channels=hp[\"hidden_dim\"],\n",
    "    num_layers=hp[\"num_layers\"],\n",
    "    out_channels=7,\n",
    "    dropout=hp[\"dropout\"],\n",
    ").to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(), lr=hp[\"lr\"], weight_decay=hp[\"l2_penalty\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tigergraph/conda/lib/python3.9/site-packages/torch/utils/tensorboard/__init__.py:4: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if not hasattr(tensorboard, '__version__') or LooseVersion(tensorboard.__version__) < LooseVersion('1.15'):\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "from pyTigerGraph.gds.metrics import Accumulator, Accuracy\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 00, Train Loss: 1.9343, Valid Loss: 1.9545, Train Accuracy: 0.0929, Valid Accuracy: 0.1020\n",
      "Epoch: 01, Train Loss: 1.7044, Valid Loss: 1.8245, Train Accuracy: 0.8000, Valid Accuracy: 0.5060\n",
      "Epoch: 02, Train Loss: 1.4190, Valid Loss: 1.6822, Train Accuracy: 0.9071, Valid Accuracy: 0.6020\n",
      "Epoch: 03, Train Loss: 1.1633, Valid Loss: 1.5156, Train Accuracy: 0.9071, Valid Accuracy: 0.6660\n",
      "Epoch: 04, Train Loss: 0.8649, Valid Loss: 1.3272, Train Accuracy: 0.9571, Valid Accuracy: 0.7000\n",
      "Epoch: 05, Train Loss: 0.6656, Valid Loss: 1.1791, Train Accuracy: 0.9571, Valid Accuracy: 0.7640\n",
      "Epoch: 06, Train Loss: 0.5014, Valid Loss: 1.0460, Train Accuracy: 0.9714, Valid Accuracy: 0.7700\n",
      "Epoch: 07, Train Loss: 0.3642, Valid Loss: 0.9679, Train Accuracy: 0.9786, Valid Accuracy: 0.7620\n",
      "Epoch: 08, Train Loss: 0.2922, Valid Loss: 0.9067, Train Accuracy: 0.9643, Valid Accuracy: 0.7460\n",
      "Epoch: 09, Train Loss: 0.2006, Valid Loss: 0.8227, Train Accuracy: 0.9786, Valid Accuracy: 0.7800\n",
      "Epoch: 10, Train Loss: 0.1874, Valid Loss: 0.8226, Train Accuracy: 0.9857, Valid Accuracy: 0.7320\n",
      "Epoch: 11, Train Loss: 0.1205, Valid Loss: 0.7543, Train Accuracy: 0.9929, Valid Accuracy: 0.7820\n",
      "Epoch: 12, Train Loss: 0.0863, Valid Loss: 0.8408, Train Accuracy: 1.0000, Valid Accuracy: 0.7420\n",
      "Epoch: 13, Train Loss: 0.0700, Valid Loss: 0.8443, Train Accuracy: 1.0000, Valid Accuracy: 0.7540\n",
      "Epoch: 14, Train Loss: 0.0654, Valid Loss: 0.7762, Train Accuracy: 0.9857, Valid Accuracy: 0.7680\n",
      "Epoch: 15, Train Loss: 0.0584, Valid Loss: 0.8157, Train Accuracy: 0.9929, Valid Accuracy: 0.7560\n",
      "Epoch: 16, Train Loss: 0.0366, Valid Loss: 0.8999, Train Accuracy: 1.0000, Valid Accuracy: 0.7540\n",
      "Epoch: 17, Train Loss: 0.0343, Valid Loss: 0.8492, Train Accuracy: 1.0000, Valid Accuracy: 0.7440\n",
      "Epoch: 18, Train Loss: 0.0246, Valid Loss: 0.8640, Train Accuracy: 1.0000, Valid Accuracy: 0.7400\n",
      "Epoch: 19, Train Loss: 0.0231, Valid Loss: 0.8484, Train Accuracy: 1.0000, Valid Accuracy: 0.7540\n"
     ]
    }
   ],
   "source": [
    "log_dir = \"logs/cora/gcn/wholegraph/\" + datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "tb_log = SummaryWriter(log_dir)\n",
    "logs = {}\n",
    "data = data.to(device)\n",
    "for epoch in range(20):\n",
    "    # Train\n",
    "    model.train()\n",
    "    acc = Accuracy()\n",
    "    # Forward pass\n",
    "    out = model(data.x.float(), data.edge_index)\n",
    "    # Calculate loss\n",
    "    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])\n",
    "    # Backward pass\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    # Evaluate\n",
    "    val_acc = Accuracy()\n",
    "    with torch.no_grad():\n",
    "        pred = out.argmax(dim=1)\n",
    "        acc.update(pred[data.train_mask], data.y[data.train_mask])\n",
    "        valid_loss = F.cross_entropy(out[data.val_mask], data.y[data.val_mask])\n",
    "        val_acc.update(pred[data.val_mask], data.y[data.val_mask])\n",
    "    # Logging\n",
    "    logs[\"loss\"] = loss.item()\n",
    "    logs[\"val_loss\"] = valid_loss.item()\n",
    "    logs[\"acc\"] = acc.value\n",
    "    logs[\"val_acc\"] = val_acc.value\n",
    "    print(\n",
    "        \"Epoch: {:02d}, Train Loss: {:.4f}, Valid Loss: {:.4f}, Train Accuracy: {:.4f}, Valid Accuracy: {:.4f}\".format(\n",
    "            epoch, logs[\"loss\"], logs[\"val_loss\"], logs[\"acc\"], logs[\"val_acc\"]\n",
    "        )\n",
    "    )\n",
    "    tb_log.add_scalars(\n",
    "        \"Loss\", {\"Train\": logs[\"loss\"], \"Validation\": logs[\"val_loss\"]}, epoch\n",
    "    )\n",
    "    tb_log.add_scalars(\n",
    "        \"Accuracy\", {\"Train\": logs[\"acc\"], \"Validation\": logs[\"val_acc\"]}, epoch\n",
    "    )\n",
    "    tb_log.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.7970\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "acc = Accuracy()\n",
    "with torch.no_grad():\n",
    "    pred = model(data.x.float(), data.edge_index).argmax(dim=1)\n",
    "    acc.update(pred[data.test_mask], data.y[data.test_mask])\n",
    "print(\"Accuracy: {:.4f}\".format(acc.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on Neighborhood Subgraphs <a name=\"train_subgraph\"></a>\n",
    "Alternatively, we train the model on the neighborhood subgraphs. Each subgraph contains the 2 hop neighborhood of certain seed vertices. This method  will allow us to train the model on graphs that are way larger than the CORA dataset because we don't load the whole graph into memory all at once. \n",
    "\n",
    "We will use the same parameters as before, but we will use the NeighborLoader to load subgraphs. Once we finish iterating over all the subgraphs generated by the loader, it is guaranteed to cover all vertices in the graph (except for those filtered by a user provided mask). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "hp = {\"batch_size\": 64, \"num_neighbors\": 10, \"num_hops\": 2, \"hidden_dim\": 64,\n",
    "      \"num_layers\": 2, \"dropout\": 0.6, \"lr\": 0.01, \"l2_penalty\": 5e-4}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct neighborhood subgraph loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we construct 3 subgraph loaders. The `train_loader` only uses vertices in the training set as seeds, the `valid_loader` only uses vertices in the validation set, and the `test_loader` only uses vertices in the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing and optimizing queries. It might take a minute if this is the first time you use this loader.\n",
      "Query installation finished.\n"
     ]
    }
   ],
   "source": [
    "train_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats=[\"x\"],\n",
    "    v_out_labels=[\"y\"],\n",
    "    v_extra_feats=[\"train_mask\",\"val_mask\",\"test_mask\"],\n",
    "    output_format=\"PyG\",\n",
    "    batch_size=hp[\"batch_size\"],\n",
    "    num_neighbors=hp[\"num_neighbors\"],\n",
    "    num_hops=hp[\"num_hops\"],\n",
    "    shuffle=True,\n",
    "    filter_by=\"train_mask\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats=[\"x\"],\n",
    "    v_out_labels=[\"y\"],\n",
    "    v_extra_feats=[\"train_mask\",\"val_mask\",\"test_mask\"],\n",
    "    output_format=\"PyG\",\n",
    "    batch_size=hp[\"batch_size\"],\n",
    "    num_neighbors=hp[\"num_neighbors\"],\n",
    "    num_hops=hp[\"num_hops\"],\n",
    "    shuffle=False,\n",
    "    filter_by=\"val_mask\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct model and optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We build a GCN model with 2 convolutional layers, and use the Adam optimizer with a learning rate of 0.01."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = GCN(\n",
    "    in_channels=1433,\n",
    "    hidden_channels=hp[\"hidden_dim\"],\n",
    "    num_layers=hp[\"num_layers\"],\n",
    "    out_channels=7,\n",
    "    dropout=hp[\"dropout\"],\n",
    ").to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(), lr=hp[\"lr\"], weight_decay=hp[\"l2_penalty\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from pyTigerGraph.gds.metrics import Accumulator, Accuracy\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, Train Batch 0, Loss 1.9169, Accuracy 0.2500\n",
      "Epoch 0, Train Batch 1, Loss 1.8491, Accuracy 0.4094\n",
      "Epoch 0, Train Batch 2, Loss 1.7908, Accuracy 0.4783\n",
      "Epoch 0, Valid Loss 1.6421, Valid Accuracy 0.6348\n",
      "Epoch 1, Train Batch 0, Loss 1.3010, Accuracy 0.8333\n",
      "Epoch 1, Train Batch 1, Loss 1.1545, Accuracy 0.8562\n",
      "Epoch 1, Train Batch 2, Loss 1.0628, Accuracy 0.8767\n",
      "Epoch 1, Valid Loss 1.3176, Valid Accuracy 0.6678\n",
      "Epoch 2, Train Batch 0, Loss 0.7232, Accuracy 0.8608\n",
      "Epoch 2, Train Batch 1, Loss 0.6530, Accuracy 0.8831\n",
      "Epoch 2, Train Batch 2, Loss 0.5754, Accuracy 0.9078\n",
      "Epoch 2, Valid Loss 1.0402, Valid Accuracy 0.7292\n",
      "Epoch 3, Train Batch 0, Loss 0.2984, Accuracy 0.9483\n",
      "Epoch 3, Train Batch 1, Loss 0.2730, Accuracy 0.9545\n",
      "Epoch 3, Train Batch 2, Loss 0.3013, Accuracy 0.9363\n",
      "Epoch 3, Valid Loss 0.8705, Valid Accuracy 0.7647\n",
      "Epoch 4, Train Batch 0, Loss 0.1879, Accuracy 0.9538\n",
      "Epoch 4, Train Batch 1, Loss 0.2048, Accuracy 0.9474\n",
      "Epoch 4, Train Batch 2, Loss 0.2044, Accuracy 0.9333\n",
      "Epoch 4, Valid Loss 0.7922, Valid Accuracy 0.7679\n",
      "Epoch 5, Train Batch 0, Loss 0.1399, Accuracy 0.9444\n",
      "Epoch 5, Train Batch 1, Loss 0.1482, Accuracy 0.9470\n",
      "Epoch 5, Train Batch 2, Loss 0.1201, Accuracy 0.9648\n",
      "Epoch 5, Valid Loss 0.8357, Valid Accuracy 0.7407\n",
      "Epoch 6, Train Batch 0, Loss 0.0688, Accuracy 0.9861\n",
      "Epoch 6, Train Batch 1, Loss 0.1047, Accuracy 0.9643\n",
      "Epoch 6, Train Batch 2, Loss 0.0981, Accuracy 0.9712\n",
      "Epoch 6, Valid Loss 0.9085, Valid Accuracy 0.7193\n",
      "Epoch 7, Train Batch 0, Loss 0.0874, Accuracy 0.9868\n",
      "Epoch 7, Train Batch 1, Loss 0.1244, Accuracy 0.9739\n",
      "Epoch 7, Train Batch 2, Loss 0.1039, Accuracy 0.9769\n",
      "Epoch 7, Valid Loss 0.9456, Valid Accuracy 0.7215\n",
      "Epoch 8, Train Batch 0, Loss 0.1709, Accuracy 0.9385\n",
      "Epoch 8, Train Batch 1, Loss 0.1048, Accuracy 0.9621\n",
      "Epoch 8, Train Batch 2, Loss 0.0794, Accuracy 0.9717\n",
      "Epoch 8, Valid Loss 0.8681, Valid Accuracy 0.7503\n",
      "Epoch 9, Train Batch 0, Loss 0.0438, Accuracy 0.9833\n",
      "Epoch 9, Train Batch 1, Loss 0.0410, Accuracy 0.9921\n",
      "Epoch 9, Train Batch 2, Loss 0.0634, Accuracy 0.9854\n",
      "Epoch 9, Valid Loss 0.8627, Valid Accuracy 0.7588\n"
     ]
    }
   ],
   "source": [
    "log_dir = \"logs/cora/gcn/subgraph/\" + datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "train_log = SummaryWriter(log_dir+\"/train\")\n",
    "valid_log = SummaryWriter(log_dir+\"/valid\")\n",
    "global_steps = 0\n",
    "logs = {}\n",
    "for epoch in range(10):\n",
    "    # Train\n",
    "    model.train()\n",
    "    epoch_train_loss = Accumulator()\n",
    "    epoch_train_acc = Accuracy()\n",
    "    # Iterate through the loader to get a stream of subgraphs instead of the whole graph\n",
    "    for bid, batch in enumerate(train_loader):\n",
    "        batchsize = batch.x.shape[0]\n",
    "        batch.to(device)\n",
    "        # Forward pass\n",
    "        out = model(batch.x.float(), batch.edge_index)\n",
    "        # Calculate loss\n",
    "        loss = F.cross_entropy(out[batch.train_mask], batch.y[batch.train_mask])\n",
    "        # Backward pass\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        epoch_train_loss.update(loss.item() * batchsize, batchsize)\n",
    "        # Predict on training data\n",
    "        with torch.no_grad():\n",
    "            pred = out.argmax(dim=1)\n",
    "            epoch_train_acc.update(pred[batch.train_mask], batch.y[batch.train_mask])\n",
    "        # Log training status after each batch\n",
    "        logs[\"loss\"] = epoch_train_loss.mean\n",
    "        logs[\"acc\"] = epoch_train_acc.value\n",
    "        print(\n",
    "            \"Epoch {}, Train Batch {}, Loss {:.4f}, Accuracy {:.4f}\".format(\n",
    "                epoch, bid, logs[\"loss\"], logs[\"acc\"]\n",
    "            )\n",
    "        )\n",
    "        train_log.add_scalar(\"Loss\", logs[\"loss\"], global_steps)\n",
    "        train_log.add_scalar(\"Accuracy\", logs[\"acc\"], global_steps)\n",
    "        train_log.flush()\n",
    "        global_steps += 1\n",
    "    # Evaluate\n",
    "    model.eval()\n",
    "    epoch_val_loss = Accumulator()\n",
    "    epoch_val_acc = Accuracy()\n",
    "    for batch in valid_loader:\n",
    "        batchsize = batch.x.shape[0]\n",
    "        batch.to(device)\n",
    "        with torch.no_grad():\n",
    "            # Forward pass\n",
    "            out = model(batch.x.float(), batch.edge_index)\n",
    "            # Calculate loss\n",
    "            valid_loss = F.cross_entropy(out[batch.val_mask], batch.y[batch.val_mask])\n",
    "            epoch_val_loss.update(valid_loss.item() * batchsize, batchsize)\n",
    "            # Prediction\n",
    "            pred = out.argmax(dim=1)\n",
    "            epoch_val_acc.update(pred[batch.val_mask], batch.y[batch.val_mask])\n",
    "    # Log testing result after each epoch\n",
    "    logs[\"val_loss\"] = epoch_val_loss.mean\n",
    "    logs[\"val_acc\"] = epoch_val_acc.value\n",
    "    print(\n",
    "        \"Epoch {}, Valid Loss {:.4f}, Valid Accuracy {:.4f}\".format(\n",
    "            epoch, logs[\"val_loss\"], logs[\"val_acc\"]\n",
    "        )\n",
    "    )\n",
    "    valid_log.add_scalar(\"Loss\", logs[\"val_loss\"], global_steps)\n",
    "    valid_log.add_scalar(\"Accuracy\", logs[\"val_acc\"], global_steps)\n",
    "    valid_log.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats=[\"x\"],\n",
    "    v_out_labels=[\"y\"],\n",
    "    v_extra_feats=[\"train_mask\",\"val_mask\",\"test_mask\"],\n",
    "    output_format=\"PyG\",\n",
    "    batch_size=hp[\"batch_size\"],\n",
    "    num_neighbors=hp[\"num_neighbors\"],\n",
    "    num_hops=hp[\"num_hops\"],\n",
    "    shuffle=False,\n",
    "    filter_by=\"test_mask\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.7610\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "acc = Accuracy()\n",
    "for batch in test_loader:\n",
    "    batch.to(device)\n",
    "    with torch.no_grad():\n",
    "        pred = model(batch.x.float(), batch.edge_index).argmax(dim=1)\n",
    "        acc.update(pred[batch.test_mask], batch.y[batch.test_mask])\n",
    "print(\"Accuracy: {:.4f}\".format(acc.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference <a name=\"inference\"></a>\n",
    "\n",
    "Finally, we use the trained model for node classification. At this stage, we typically do inference/prediction for specific nodes instead of random batches, so we will create a new data loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats=[\"x\"],\n",
    "    v_out_labels=[\"y\"],\n",
    "    v_extra_feats=[\"train_mask\",\"val_mask\",\"test_mask\"],\n",
    "    output_format=\"PyG\",\n",
    "    num_neighbors=hp[\"num_neighbors\"],\n",
    "    num_hops=hp[\"num_hops\"],\n",
    "    shuffle=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch specific nodes by their IDs and do prediction. \n",
    "# Each node is represented by a dict with two mandatory keys: primary_id and type.\n",
    "input_nodes = [{\"primary_id\": 7, \"type\": \"Paper\"}, \n",
    "               {\"primary_id\": 999, \"type\": \"Paper\"}]\n",
    "data = infer_loader.fetch(input_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 16], x=[15, 1433], y=[15], train_mask=[15], val_mask=[15], test_mask=[15], is_seed=[15], primary_id=[15])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The returned data are the neighborhood subgraphs of the input nodes.\n",
    "# The original IDs of the nodes in the subgraphs are stored in the \n",
    "# `primary_id` attribute.\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ID: Label\n",
      "7:3\n",
      "999:2\n",
      "1358:2\n",
      "346:2\n",
      "208:3\n",
      "754:2\n",
      "823:2\n",
      "1708:2\n",
      "72:2\n",
      "1646:2\n",
      "580:2\n",
      "1568:2\n",
      "1747:2\n",
      "902:2\n",
      "524:2\n"
     ]
    }
   ],
   "source": [
    "# Predict. Predictions for both the input nodes and others in their \n",
    "# neighborhoods are generated.\n",
    "model.eval()\n",
    "pred = model(data.x.float(), data.edge_index).argmax(dim=1)\n",
    "print(\"ID: Label\")\n",
    "for i,j in zip(data.primary_id, pred):\n",
    "    print(\"{}:{}\".format(i, j.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "pytorch-gpu.1-9.m81",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-9:m81"
  },
  "kernelspec": {
   "display_name": "Python 3.8.9 64-bit",
   "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.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}