GNNs/DGL/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. [DGL](https://www.dgl.ai/)'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 dgl -f https://data.dgl.ai/wheels/repo.html\n",
    "#!pip install psutil # Required for DGL\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": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb64eaa96c72469893d6931cd5535abf",
       "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": 3,
   "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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cf23cb84fc904f5a9317aca8590185f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "CytoscapeWidget(cytoscape_layout={'name': 'circle', 'animate': True, 'padding': 1}, cytoscape_style=[{'selecto…"
      ]
     },
     "execution_count": 4,
     "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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Paper': 2708}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conn.getVertexCount('*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Cite': 10556}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conn.getEdgeCount()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train/validation/test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 notebook 1_data_processing 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": null,
   "metadata": {},
   "outputs": [],
   "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",
    "\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "hp = {\"hidden_dim\": 64, \n",
    "      \"num_layers\": 2, \n",
    "      \"dropout\": 0.6, \n",
    "      \"lr\": 0.01, \n",
    "      \"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 for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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=\"DGL\",\n",
    "    shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the whole graph from the loader in DGL format\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl.function as fn\n",
    "import dgl.nn.pytorch as dglnn\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GCN, self).__init__()\n",
    "        self.layer1 = dglnn.conv.GraphConv(1433, hp['hidden_dim'])\n",
    "        self.layer2 = dglnn.conv.GraphConv(hp['hidden_dim'], 7)\n",
    "\n",
    "    def forward(self, g, features):\n",
    "        x = F.relu(self.layer1(g, features))\n",
    "        x = self.layer2(g, x)\n",
    "        return x\n",
    "\n",
    "model = GCN()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(), lr=hp[\"lr\"], weight_decay=hp[\"l2_penalty\"]\n",
    ")\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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, data.ndata[\"x\"])\n",
    "    # Calculate loss\n",
    "    loss = F.cross_entropy(out[data.ndata[\"train_mask\"]], data.ndata[\"y\"][data.ndata[\"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.ndata[\"train_mask\"]], data.ndata[\"y\"][data.ndata[\"train_mask\"]])\n",
    "        valid_loss = F.cross_entropy(out[data.ndata[\"val_mask\"]], data.ndata[\"y\"][data.ndata[\"val_mask\"]])\n",
    "        val_acc.update(pred[data.ndata[\"val_mask\"]], data.ndata[\"y\"][data.ndata[\"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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "acc = Accuracy()\n",
    "with torch.no_grad():\n",
    "    pred = model(data, data.ndata[\"x\"]).argmax(dim=1)\n",
    "    acc.update(pred[data.ndata[\"test_mask\"]], data.ndata[\"y\"][data.ndata[\"test_mask\"]])\n",
    "print(\"Accuracy: {:.4f}\".format(acc.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on Neighborhood Subgraphs <a name=\"train_subgraph\"></a>\n",
    "\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "hp = {\"batch_size\": 64, \n",
    "      \"num_neighbors\": 10, \n",
    "      \"num_hops\": 2, \n",
    "      \"hidden_dim\": 64, \n",
    "      \"num_layers\": 2, \n",
    "      \"dropout\": 0.6, \n",
    "      \"lr\": 0.01, \n",
    "      \"l2_penalty\": 5e-4}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct neighborhood subgraph loader\n",
    "\n",
    "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": null,
   "metadata": {},
   "outputs": [],
   "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=\"DGL\",\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",
    "    add_self_loop=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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=\"DGL\",\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",
    "    add_self_loop=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct model and optimizer\n",
    "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = GCN().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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "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",
    "    for bid, batch in enumerate(train_loader):\n",
    "        batchsize = batch.num_nodes()\n",
    "        batch.to(device)\n",
    "        # Forward pass\n",
    "        out = model(batch, batch.ndata[\"x\"])\n",
    "        # Calculate loss\n",
    "        loss = F.cross_entropy(out[batch.ndata[\"is_seed\"]], batch.ndata[\"y\"][batch.ndata[\"is_seed\"]])\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.ndata[\"is_seed\"]], batch.ndata[\"y\"][batch.ndata[\"is_seed\"]])\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.num_nodes()\n",
    "        batch.to(device)\n",
    "        with torch.no_grad():\n",
    "            # Forward pass\n",
    "            out = model(batch, batch.ndata[\"x\"])\n",
    "            # Calculate loss\n",
    "            valid_loss = F.cross_entropy(out[batch.ndata[\"val_mask\"]], batch.ndata[\"y\"][batch.ndata[\"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.ndata[\"val_mask\"]], batch.ndata[\"y\"][batch.ndata[\"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": null,
   "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=\"DGL\",\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",
    "    add_self_loop=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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, batch.ndata[\"x\"]).argmax(dim=1)\n",
    "        acc.update(pred[batch.ndata[\"test_mask\"]], batch.ndata[\"y\"][batch.ndata[\"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": null,
   "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=\"DGL\",\n",
    "    num_neighbors=hp[\"num_neighbors\"],\n",
    "    num_hops=hp[\"num_hops\"],\n",
    "    shuffle=False,\n",
    "    add_self_loop=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict. Predictions for both the input nodes and others in their \n",
    "# neighborhoods are generated.\n",
    "model.eval()\n",
    "pred = model(data, data.ndata[\"x\"]).argmax(dim=1)\n",
    "print(\"ID: Label\")\n",
    "for i,j in zip(data.extra_data[\"primary_id\"], pred):\n",
    "    print(\"{}:{}\".format(i, j.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch",
   "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
}