GNNs/PyG/hgat_node_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Heterogeneous Graph Attention Network\n",
    "This notebook demonstrates the training of Heterogeneous Graph Attention Networks (HGAT) with TigerGraph ML Workbench. [Pytorch Geometric](https://pytorch-geometric.readthedocs.io)'s implementation of HGAT is used here. We train the model on the IMDB dataset from [PyG datasets](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.IMDB) with TigerGraph as the data store. The dataset contains 3 types of vertices: 4278 movies, 5257 actors, and 2081 directors; and 4 types of edges: 12828 actor to movie edges, 12828 movie to actor edges, 4278 director to movie edges, and 4278 movie to director edges. Each vertex is described by a 0/1-valued word vector indicating the absence/presence of the corresponding keywords from the plot (for movie) or from movies they participated (for actors and directors). Each movie is classified into one of three classes, action, comedy, and drama according to their genre. The goal is to predict the class of each movie in the graph.\n",
    "\n",
    "The following libraries are required to run this notebook. Uncomment to install them if necessary. You 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": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecbfd9b8946341b3872ca12dbd6f3613",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0/441353 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from pyTigerGraph.datasets import Datasets\n",
    "\n",
    "dataset = Datasets(\"imdb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Checking database ----\n",
      "A graph with name imdb 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": "2060504b36f04ecdbdbcf473bf7b19cc",
       "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": [
       "{'Movie': 4278, 'Actor': 5257, 'Director': 2081}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conn.getVertexCount('*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'actor_movie': 12828,\n",
       " 'director_movie': 4278,\n",
       " 'movie_actor': 12828,\n",
       " 'movie_director': 4278}"
      ]
     },
     "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": 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 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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of movies in training set: 400\n",
      "Number of movies in validation set: 400\n",
      "Number of movies in test set: 3478\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Number of movies in training set:\",\n",
    "    conn.getVertexCount(\"Movie\", where=\"train_mask!=0\"),\n",
    ")\n",
    "print(\n",
    "    \"Number of movies in validation set:\",\n",
    "    conn.getVertexCount(\"Movie\", where=\"val_mask!=0\"),\n",
    ")\n",
    "print(\n",
    "    \"Number of movies in test set:\", \n",
    "    conn.getVertexCount(\"Movie\", 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": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "hp = {\n",
    "    \"num_heads\": 2,\n",
    "    \"hidden_dim\": 8,\n",
    "    \"num_layers\": 2,\n",
    "    \"dropout\": 0.1,\n",
    "    \"lr\": 0.01,\n",
    "    \"l2_penalty\": 0.0001,\n",
    "}"
   ]
  },
  {
   "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": 9,
   "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={\"Movie\": [\"x\"], \"Actor\": [\"x\"], \"Director\": [\"x\"]}, \n",
    "    v_out_labels={\"Movie\": [\"y\"]},\n",
    "    v_extra_feats={\"Movie\": [\"train_mask\", \"val_mask\", \"test_mask\"]},\n",
    "    num_batches=1,\n",
    "    output_format=\"PyG\",\n",
    "    shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  \u001b[1mMovie\u001b[0m={\n",
       "    x=[4278, 3066],\n",
       "    y=[4278],\n",
       "    train_mask=[4278],\n",
       "    val_mask=[4278],\n",
       "    test_mask=[4278]\n",
       "  },\n",
       "  \u001b[1mActor\u001b[0m={ x=[5257, 3066] },\n",
       "  \u001b[1mDirector\u001b[0m={ x=[2081, 3066] },\n",
       "  \u001b[1m(Movie, movie_actor, Actor)\u001b[0m={ edge_index=[2, 12828] },\n",
       "  \u001b[1m(Movie, movie_director, Director)\u001b[0m={ edge_index=[2, 4278] },\n",
       "  \u001b[1m(Actor, actor_movie, Movie)\u001b[0m={ edge_index=[2, 12828] },\n",
       "  \u001b[1m(Director, director_movie, Movie)\u001b[0m={ edge_index=[2, 4278] }\n",
       ")"
      ]
     },
     "execution_count": 10,
     "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 GAT model with 2 convolutional layers, and then convert it to a heterogenous GAT model. We use the Adam optimizer with a learning rate of 0.01 to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GATConv, to_hetero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Create a normal (homogeneous) GAT model\n",
    "class GAT(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self, num_features, num_layers, out_dim, dropout, hidden_dim, num_heads\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.dropout = dropout\n",
    "        self.layers = torch.nn.ModuleList()\n",
    "        for i in range(num_layers):\n",
    "            in_units = num_features if i == 0 else hidden_dim * num_heads\n",
    "            out_units = out_dim if i == (num_layers - 1) else hidden_dim\n",
    "            heads = 1 if i == (num_layers - 1) else num_heads\n",
    "            self.layers.append(\n",
    "                GATConv(in_units, out_units, heads=heads, dropout=dropout, add_self_loops=False)\n",
    "            )\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        for layer in self.layers:\n",
    "            layer.reset_parameters()\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = x.float()\n",
    "        for layer in self.layers[:-1]:\n",
    "            x = layer(x, edge_index)\n",
    "            x = F.elu(x)\n",
    "            x = F.dropout(x, p=self.dropout, training=self.training)\n",
    "        x = self.layers[-1](x, edge_index)\n",
    "        return x\n",
    "    \n",
    "model = GAT(\n",
    "    num_features=3066,\n",
    "    num_layers=hp[\"num_layers\"],\n",
    "    out_dim=3,\n",
    "    dropout=hp[\"dropout\"],\n",
    "    hidden_dim=hp[\"hidden_dim\"],\n",
    "    num_heads=hp[\"num_heads\"],\n",
    ")\n",
    "\n",
    "# Convert it to a heterogeneous model. See https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.to_hetero_transformer.to_hetero for details.\n",
    "model = to_hetero(model, data.metadata(), aggr='sum').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": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/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(\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": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 00, Train Loss: 1.0999, Valid Loss: 1.0952, Train Accuracy: 0.3200, Valid Accuracy: 0.3600\n",
      "Epoch: 01, Train Loss: 0.9461, Valid Loss: 1.0454, Train Accuracy: 0.7500, Valid Accuracy: 0.4825\n",
      "Epoch: 02, Train Loss: 0.8101, Valid Loss: 1.0224, Train Accuracy: 0.8100, Valid Accuracy: 0.4750\n",
      "Epoch: 03, Train Loss: 0.6942, Valid Loss: 1.0109, Train Accuracy: 0.8475, Valid Accuracy: 0.4600\n",
      "Epoch: 04, Train Loss: 0.5879, Valid Loss: 1.0056, Train Accuracy: 0.9125, Valid Accuracy: 0.4575\n",
      "Epoch: 05, Train Loss: 0.4926, Valid Loss: 0.9867, Train Accuracy: 0.9525, Valid Accuracy: 0.5075\n",
      "Epoch: 06, Train Loss: 0.3987, Valid Loss: 1.0070, Train Accuracy: 0.9700, Valid Accuracy: 0.4775\n",
      "Epoch: 07, Train Loss: 0.3256, Valid Loss: 0.9922, Train Accuracy: 0.9900, Valid Accuracy: 0.5200\n",
      "Epoch: 08, Train Loss: 0.2591, Valid Loss: 1.0082, Train Accuracy: 0.9875, Valid Accuracy: 0.5000\n",
      "Epoch: 09, Train Loss: 0.2147, Valid Loss: 1.0477, Train Accuracy: 0.9950, Valid Accuracy: 0.4900\n",
      "Epoch: 10, Train Loss: 0.1757, Valid Loss: 1.0885, Train Accuracy: 0.9975, Valid Accuracy: 0.5425\n",
      "Epoch: 11, Train Loss: 0.1395, Valid Loss: 1.1011, Train Accuracy: 0.9900, Valid Accuracy: 0.5325\n",
      "Epoch: 12, Train Loss: 0.1192, Valid Loss: 1.1494, Train Accuracy: 0.9950, Valid Accuracy: 0.5175\n",
      "Epoch: 13, Train Loss: 0.0942, Valid Loss: 1.1812, Train Accuracy: 0.9925, Valid Accuracy: 0.5275\n",
      "Epoch: 14, Train Loss: 0.0872, Valid Loss: 1.2335, Train Accuracy: 0.9975, Valid Accuracy: 0.5050\n",
      "Epoch: 15, Train Loss: 0.0648, Valid Loss: 1.2654, Train Accuracy: 1.0000, Valid Accuracy: 0.5100\n",
      "Epoch: 16, Train Loss: 0.0524, Valid Loss: 1.2933, Train Accuracy: 0.9925, Valid Accuracy: 0.5300\n",
      "Epoch: 17, Train Loss: 0.0438, Valid Loss: 1.3760, Train Accuracy: 0.9975, Valid Accuracy: 0.5350\n",
      "Epoch: 18, Train Loss: 0.0364, Valid Loss: 1.3579, Train Accuracy: 1.0000, Valid Accuracy: 0.5225\n",
      "Epoch: 19, Train Loss: 0.0345, Valid Loss: 1.4159, Train Accuracy: 1.0000, Valid Accuracy: 0.5350\n"
     ]
    }
   ],
   "source": [
    "log_dir = \"logs/imdb/hgat/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_dict, data.edge_index_dict)\n",
    "    # Calculate loss on movie vertices in the training set only\n",
    "    mask = data['Movie'].train_mask\n",
    "    loss = F.cross_entropy(out[\"Movie\"][mask], data[\"Movie\"].y[mask])\n",
    "    # Backward pass\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    # Update model\n",
    "    optimizer.step()\n",
    "    # Evaluate\n",
    "    val_acc = Accuracy()\n",
    "    with torch.no_grad():\n",
    "        pred = out['Movie'].argmax(dim=1)\n",
    "        acc.update(pred[mask], data['Movie'].y[mask])\n",
    "        mask = data['Movie'].val_mask\n",
    "        valid_loss = F.cross_entropy(out['Movie'][mask], data['Movie'].y[mask])\n",
    "        val_acc.update(pred[mask], data['Movie'].y[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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.5190\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "acc = Accuracy()\n",
    "with torch.no_grad():\n",
    "    pred = model(data.x_dict, data.edge_index_dict)[\"Movie\"].argmax(dim=1)\n",
    "    mask = data[\"Movie\"].test_mask\n",
    "    acc.update(pred[mask], data[\"Movie\"].y[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 IMDB 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": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "hp = {\n",
    "    \"num_heads\": 2,\n",
    "    \"hidden_dim\": 8,\n",
    "    \"num_layers\": 2,\n",
    "    \"dropout\": 0.2,\n",
    "    \"lr\": 0.01,\n",
    "    \"l2_penalty\": 0.0001,\n",
    "    \"batch_size\": 128, \n",
    "    \"num_neighbors\": 10, \n",
    "    \"num_hops\": 2\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": 18,
   "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={\"Movie\": [\"x\"], \"Actor\": [\"x\"], \"Director\": [\"x\"]}, \n",
    "    v_out_labels={\"Movie\": [\"y\"]},\n",
    "    v_extra_feats={\"Movie\": [\"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={\"Movie\":\"train_mask\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats={\"Movie\": [\"x\"], \"Actor\": [\"x\"], \"Director\": [\"x\"]}, \n",
    "    v_out_labels={\"Movie\": [\"y\"]},\n",
    "    v_extra_feats={\"Movie\": [\"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={\"Movie\":\"val_mask\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct model and optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We build a GAT model with 2 convolutional layers, and then convert it to a heterogenous GAT model. We use the Adam optimizer with a learning rate of 0.01 to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GATConv, to_hetero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Create a normal (homogeneous) GAT model\n",
    "class GAT(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self, num_features, num_layers, out_dim, dropout, hidden_dim, num_heads\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.dropout = dropout\n",
    "        self.layers = torch.nn.ModuleList()\n",
    "        for i in range(num_layers):\n",
    "            in_units = num_features if i == 0 else hidden_dim * num_heads\n",
    "            out_units = out_dim if i == (num_layers - 1) else hidden_dim\n",
    "            heads = 1 if i == (num_layers - 1) else num_heads\n",
    "            self.layers.append(\n",
    "                GATConv(in_units, out_units, heads=heads, dropout=dropout, add_self_loops=False)\n",
    "            )\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        for layer in self.layers:\n",
    "            layer.reset_parameters()\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = x.float()\n",
    "        for layer in self.layers[:-1]:\n",
    "            x = layer(x, edge_index)\n",
    "            x = F.elu(x)\n",
    "            x = F.dropout(x, p=self.dropout, training=self.training)\n",
    "        x = self.layers[-1](x, edge_index)\n",
    "        return x\n",
    "    \n",
    "model = GAT(\n",
    "    num_features=3066,\n",
    "    num_layers=hp[\"num_layers\"],\n",
    "    out_dim=3,\n",
    "    dropout=hp[\"dropout\"],\n",
    "    hidden_dim=hp[\"hidden_dim\"],\n",
    "    num_heads=hp[\"num_heads\"],\n",
    ")\n",
    "\n",
    "# Convert it to a heterogeneous model. See https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.to_hetero_transformer.to_hetero for details.\n",
    "metadata = (['Actor', 'Movie', 'Director'], \n",
    "            [('Actor', 'actor_movie', 'Movie'), \n",
    "             ('Movie', 'movie_actor', 'Actor'), \n",
    "             ('Movie', 'movie_director', 'Director'), \n",
    "             ('Director', 'director_movie', 'Movie')])\n",
    "model = to_hetero(model, metadata, aggr='sum').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": 22,
   "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": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, Train Batch 0, Loss 1.1123, Accuracy 0.3086\n",
      "Epoch 0, Train Batch 1, Loss 1.0735, Accuracy 0.4022\n",
      "Epoch 0, Train Batch 2, Loss 1.0822, Accuracy 0.3693\n",
      "Epoch 0, Train Batch 3, Loss 1.0856, Accuracy 0.3675\n",
      "Epoch 0, Valid Loss 1.0530, Valid Accuracy 0.3925\n",
      "Epoch 1, Train Batch 0, Loss 0.9177, Accuracy 0.5484\n",
      "Epoch 1, Train Batch 1, Loss 0.9169, Accuracy 0.5459\n",
      "Epoch 1, Train Batch 2, Loss 0.8784, Accuracy 0.5973\n",
      "Epoch 1, Train Batch 3, Loss 0.8585, Accuracy 0.6275\n",
      "Epoch 1, Valid Loss 1.0094, Valid Accuracy 0.4650\n",
      "Epoch 2, Train Batch 0, Loss 0.7882, Accuracy 0.7553\n",
      "Epoch 2, Train Batch 1, Loss 0.7354, Accuracy 0.7730\n",
      "Epoch 2, Train Batch 2, Loss 0.7057, Accuracy 0.7973\n",
      "Epoch 2, Train Batch 3, Loss 0.6940, Accuracy 0.8050\n",
      "Epoch 2, Valid Loss 1.0000, Valid Accuracy 0.5225\n",
      "Epoch 3, Train Batch 0, Loss 0.6163, Accuracy 0.8454\n",
      "Epoch 3, Train Batch 1, Loss 0.5857, Accuracy 0.8357\n",
      "Epoch 3, Train Batch 2, Loss 0.5866, Accuracy 0.8384\n",
      "Epoch 3, Train Batch 3, Loss 0.5947, Accuracy 0.8200\n",
      "Epoch 3, Valid Loss 1.0239, Valid Accuracy 0.4875\n",
      "Epoch 4, Train Batch 0, Loss 0.5037, Accuracy 0.7946\n",
      "Epoch 4, Train Batch 1, Loss 0.4738, Accuracy 0.8318\n",
      "Epoch 4, Train Batch 2, Loss 0.4801, Accuracy 0.8449\n",
      "Epoch 4, Train Batch 3, Loss 0.4578, Accuracy 0.8700\n",
      "Epoch 4, Valid Loss 0.9940, Valid Accuracy 0.5150\n",
      "Epoch 5, Train Batch 0, Loss 0.4366, Accuracy 0.8710\n",
      "Epoch 5, Train Batch 1, Loss 0.4044, Accuracy 0.8860\n",
      "Epoch 5, Train Batch 2, Loss 0.4181, Accuracy 0.8618\n",
      "Epoch 5, Train Batch 3, Loss 0.4208, Accuracy 0.8600\n",
      "Epoch 5, Valid Loss 1.0900, Valid Accuracy 0.4875\n",
      "Epoch 6, Train Batch 0, Loss 0.3842, Accuracy 0.8558\n",
      "Epoch 6, Train Batch 1, Loss 0.3925, Accuracy 0.8502\n",
      "Epoch 6, Train Batch 2, Loss 0.3591, Accuracy 0.8677\n",
      "Epoch 6, Train Batch 3, Loss 0.3549, Accuracy 0.8600\n",
      "Epoch 6, Valid Loss 1.1037, Valid Accuracy 0.5025\n",
      "Epoch 7, Train Batch 0, Loss 0.2445, Accuracy 0.9316\n",
      "Epoch 7, Train Batch 1, Loss 0.2464, Accuracy 0.9265\n",
      "Epoch 7, Train Batch 2, Loss 0.2652, Accuracy 0.9288\n",
      "Epoch 7, Train Batch 3, Loss 0.2702, Accuracy 0.9225\n",
      "Epoch 7, Valid Loss 1.1855, Valid Accuracy 0.5000\n",
      "Epoch 8, Train Batch 0, Loss 0.3005, Accuracy 0.8962\n",
      "Epoch 8, Train Batch 1, Loss 0.2727, Accuracy 0.9100\n",
      "Epoch 8, Train Batch 2, Loss 0.2487, Accuracy 0.9133\n",
      "Epoch 8, Train Batch 3, Loss 0.2633, Accuracy 0.9050\n",
      "Epoch 8, Valid Loss 1.2499, Valid Accuracy 0.4725\n",
      "Epoch 9, Train Batch 0, Loss 0.2957, Accuracy 0.8889\n",
      "Epoch 9, Train Batch 1, Loss 0.2826, Accuracy 0.8950\n",
      "Epoch 9, Train Batch 2, Loss 0.2629, Accuracy 0.9027\n",
      "Epoch 9, Train Batch 3, Loss 0.2606, Accuracy 0.9075\n",
      "Epoch 9, Valid Loss 1.2367, Valid Accuracy 0.4950\n"
     ]
    }
   ],
   "source": [
    "log_dir = \"logs/imdb/hgat/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[\"Movie\"].x.shape[0]\n",
    "        batch.to(device)\n",
    "        # Forward pass\n",
    "        out = model(batch.x_dict, batch.edge_index_dict)\n",
    "        # Calculate loss\n",
    "        mask = batch[\"Movie\"].is_seed\n",
    "        loss = F.cross_entropy(out[\"Movie\"][mask], batch[\"Movie\"].y[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[\"Movie\"].argmax(dim=1)\n",
    "            epoch_train_acc.update(pred[mask], batch[\"Movie\"].y[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[\"Movie\"].x.shape[0]\n",
    "        batch.to(device)\n",
    "        with torch.no_grad():\n",
    "            # Forward pass\n",
    "            out = model(batch.x_dict, batch.edge_index_dict)\n",
    "            # Calculate loss\n",
    "            mask = batch[\"Movie\"].is_seed\n",
    "            valid_loss = F.cross_entropy(out[\"Movie\"][mask], batch[\"Movie\"].y[mask])\n",
    "            epoch_val_loss.update(valid_loss.item() * batchsize, batchsize)\n",
    "            # Prediction\n",
    "            pred = out[\"Movie\"].argmax(dim=1)\n",
    "            epoch_val_acc.update(pred[mask], batch[\"Movie\"].y[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": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats={\"Movie\": [\"x\"], \"Actor\": [\"x\"], \"Director\": [\"x\"]}, \n",
    "    v_out_labels={\"Movie\": [\"y\"]},\n",
    "    v_extra_feats={\"Movie\": [\"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={\"Movie\":\"test_mask\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.4888\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_dict, batch.edge_index_dict)[\"Movie\"].argmax(dim=1)\n",
    "        mask = batch[\"Movie\"].is_seed\n",
    "        acc.update(pred[mask], batch[\"Movie\"].y[mask])\n",
    "print(\"Accuracy: {:.4f}\".format(acc.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "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": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_loader = conn.gds.neighborLoader(\n",
    "    v_in_feats={\"Movie\": [\"x\"], \"Actor\": [\"x\"], \"Director\": [\"x\"]}, \n",
    "    v_out_labels={\"Movie\": [\"y\"]},\n",
    "    v_extra_feats={\"Movie\": [\"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": 28,
   "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\": \"Movie\"}, \n",
    "               {\"primary_id\": 55, \"type\": \"Movie\"}]\n",
    "data = infer_loader.fetch(input_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  \u001b[1mMovie\u001b[0m={\n",
       "    x=[53, 3066],\n",
       "    y=[53],\n",
       "    train_mask=[53],\n",
       "    val_mask=[53],\n",
       "    test_mask=[53],\n",
       "    is_seed=[53],\n",
       "    primary_id=[53]\n",
       "  },\n",
       "  \u001b[1mActor\u001b[0m={\n",
       "    x=[6, 3066],\n",
       "    is_seed=[6],\n",
       "    primary_id=[6]\n",
       "  },\n",
       "  \u001b[1mDirector\u001b[0m={\n",
       "    x=[2, 3066],\n",
       "    is_seed=[2],\n",
       "    primary_id=[2]\n",
       "  },\n",
       "  \u001b[1m(Movie, movie_actor, Actor)\u001b[0m={ edge_index=[2, 6] },\n",
       "  \u001b[1m(Movie, movie_director, Director)\u001b[0m={ edge_index=[2, 2] },\n",
       "  \u001b[1m(Actor, actor_movie, Movie)\u001b[0m={ edge_index=[2, 54] },\n",
       "  \u001b[1m(Director, director_movie, Movie)\u001b[0m={ edge_index=[2, 11] }\n",
       ")"
      ]
     },
     "execution_count": 29,
     "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": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ID: Label\n",
      "7:0\n",
      "55:0\n",
      "1543:0\n",
      "963:0\n",
      "211:2\n",
      "178:0\n",
      "1074:0\n",
      "387:0\n",
      "553:0\n",
      "856:0\n",
      "1327:0\n",
      "2184:0\n",
      "2897:0\n",
      "109:0\n",
      "712:0\n",
      "3597:0\n",
      "1153:0\n",
      "3137:0\n",
      "4157:0\n",
      "1930:0\n",
      "1863:0\n",
      "1901:0\n",
      "3983:0\n",
      "633:0\n",
      "3150:0\n",
      "76:0\n",
      "9:0\n",
      "1899:0\n",
      "2025:0\n",
      "3718:0\n",
      "111:0\n",
      "1433:0\n",
      "138:0\n",
      "2077:0\n",
      "22:0\n",
      "40:0\n",
      "15:0\n",
      "2789:2\n",
      "3124:0\n",
      "3413:0\n",
      "1346:2\n",
      "464:0\n",
      "89:0\n",
      "1530:0\n",
      "2263:0\n",
      "659:0\n",
      "326:0\n",
      "3454:2\n",
      "3174:0\n",
      "70:0\n",
      "520:0\n",
      "1206:0\n",
      "2382:0\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_dict, data.edge_index_dict)[\"Movie\"].argmax(dim=1)\n",
    "print(\"ID: Label\")\n",
    "for i,j in zip(data[\"Movie\"].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
}