{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyMWE8zjXNzesRPRk8RK7jsv"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Node Classification\n",
        "This tutorial shows how to train a multi-layer GraphSAGE for node\n",
        "classification on ``ogbn-arxiv`` provided by [Open Graph\n",
        "Benchmark (OGB)](https://ogb.stanford.edu/). The dataset contains around\n",
        "170 thousand nodes and 1 million edges.\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dmlc/dgl/blob/master/notebooks/stochastic_training/node_classification.ipynb) [![GitHub](https://img.shields.io/badge/-View%20on%20GitHub-181717?logo=github&logoColor=ffffff)](https://github.com/dmlc/dgl/blob/master/notebooks/stochastic_training/node_classification.ipynb)\n",
        "\n",
        "By the end of this tutorial, you will be able to\n",
        "\n",
        "-  Train a GNN model for node classification on a single GPU with DGL's\n",
        "   neighbor sampling components."
      ],
      "metadata": {
        "id": "OxbY2KlG4ZfJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install DGL package"
      ],
      "metadata": {
        "id": "mzZKrVVk6Y_8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Install required packages.\n",
        "import os\n",
        "import torch\n",
        "import numpy as np\n",
        "os.environ['TORCH'] = torch.__version__\n",
        "os.environ['DGLBACKEND'] = \"pytorch\"\n",
        "\n",
        "# Install the CPU version.\n",
        "device = torch.device(\"cpu\")\n",
        "!pip install --pre dgl -f https://data.dgl.ai/wheels-test/repo.html\n",
        "\n",
        "try:\n",
        "    import dgl\n",
        "    import dgl.graphbolt as gb\n",
        "    installed = True\n",
        "except ImportError as error:\n",
        "    installed = False\n",
        "    print(error)\n",
        "print(\"DGL installed!\" if installed else \"DGL not found!\")"
      ],
      "metadata": {
        "id": "QcpjTazg6hEo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading Dataset\n",
        "`ogbn-arxiv` is already prepared as ``BuiltinDataset`` in **GraphBolt**."
      ],
      "metadata": {
        "id": "XWdRZAM-51Cb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = gb.BuiltinDataset(\"ogbn-arxiv\").load()"
      ],
      "metadata": {
        "id": "RnJkkSKhWiUG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Dataset consists of graph, feature and tasks. You can get the training-validation-test set from the tasks. Seed nodes and corresponding labels are already stored in each training-validation-test set. Other metadata such as number of classes are also stored in the tasks. In this dataset, there is only one task: `node classification`."
      ],
      "metadata": {
        "id": "S8avoKBiXA9j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "graph = dataset.graph\n",
        "feature = dataset.feature\n",
        "train_set = dataset.tasks[0].train_set\n",
        "valid_set = dataset.tasks[0].validation_set\n",
        "test_set = dataset.tasks[0].test_set\n",
        "task_name = dataset.tasks[0].metadata[\"name\"]\n",
        "num_classes = dataset.tasks[0].metadata[\"num_classes\"]\n",
        "print(f\"Task: {task_name}. Number of classes: {num_classes}\")"
      ],
      "metadata": {
        "id": "IXGZmgIaXJWQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## How DGL Handles Computation Dependency¶\n",
        "The computation dependency for message passing of a single node can be described as a series of message flow graphs (MFG).\n",
        "\n",
        "![DGL Computation](https://data.dgl.ai/tutorial/img/bipartite.gif)"
      ],
      "metadata": {
        "id": "y8yn77Kg6HkW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defining Neighbor Sampler and Data Loader in DGL\n",
        "\n",
        "DGL provides tools to iterate over the dataset in minibatches while generating the computation dependencies to compute their outputs with the MFGs above. For node classification, you can use `dgl.graphbolt.MultiProcessDataLoader` for iterating over the dataset. It accepts a data pipe that generates minibatches of nodes and their labels, sample neighbors for each node, and generate the computation dependencies in the form of MFGs. Feature fetching, block creation and copying to target device are also supported. All these operations are split into separate stages in the data pipe, so that you can customize the data pipeline by inserting your own operations.\n",
        "\n",
        "Let’s say that each node will gather messages from 4 neighbors on each layer. The code defining the data loader and neighbor sampler will look like the following.\n"
      ],
      "metadata": {
        "id": "q7GrcJTnZQjt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "datapipe = gb.ItemSampler(train_set, batch_size=1024, shuffle=True)\n",
        "datapipe = datapipe.sample_neighbor(graph, [4, 4])\n",
        "datapipe = datapipe.fetch_feature(feature, node_feature_keys=[\"feat\"])\n",
        "datapipe = datapipe.to_dgl()\n",
        "datapipe = datapipe.copy_to(device)\n",
        "train_dataloader = gb.MultiProcessDataLoader(datapipe, num_workers=0)"
      ],
      "metadata": {
        "id": "yQVYDO0ZbBvi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can iterate over the data loader and a `DGLMiniBatch` object is yielded.\n",
        "\n"
      ],
      "metadata": {
        "id": "7Rp12SUhbEV1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data = next(iter(train_dataloader))\n",
        "print(data)"
      ],
      "metadata": {
        "id": "V7vQiKj2bL_o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can get the input node IDs from MFGs."
      ],
      "metadata": {
        "id": "-eBuPnT-bS-o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "mfgs = data.blocks\n",
        "input_nodes = mfgs[0].srcdata[dgl.NID]\n",
        "print(f\"Input nodes: {input_nodes}.\")"
      ],
      "metadata": {
        "id": "bN4sgZqFbUvd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defining Model\n",
        "Let’s consider training a 2-layer GraphSAGE with neighbor sampling. The model can be written as follows:\n",
        "\n"
      ],
      "metadata": {
        "id": "fV6epnRxbZl4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from dgl.nn import SAGEConv\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "    def __init__(self, in_feats, h_feats, num_classes):\n",
        "        super(Model, self).__init__()\n",
        "        self.conv1 = SAGEConv(in_feats, h_feats, aggregator_type=\"mean\")\n",
        "        self.conv2 = SAGEConv(h_feats, num_classes, aggregator_type=\"mean\")\n",
        "        self.h_feats = h_feats\n",
        "\n",
        "    def forward(self, mfgs, x):\n",
        "        h = self.conv1(mfgs[0], x)\n",
        "        h = F.relu(h)\n",
        "        h = self.conv2(mfgs[1], h)\n",
        "        return h\n",
        "\n",
        "\n",
        "in_size = feature.size(\"node\", None, \"feat\")[0]\n",
        "model = Model(in_size, 64, num_classes).to(device)"
      ],
      "metadata": {
        "id": "iKhEIL0Ccmwx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defining Training Loop\n",
        "\n",
        "The following initializes the model and defines the optimizer.\n"
      ],
      "metadata": {
        "id": "OGLN3kCcwCA8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "opt = torch.optim.Adam(model.parameters())"
      ],
      "metadata": {
        "id": "dET8i_hewLUi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "When computing the validation score for model selection, usually you can also do neighbor sampling. To do that, you need to define another data loader."
      ],
      "metadata": {
        "id": "leZvFP4GwMcq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "datapipe = gb.ItemSampler(valid_set, batch_size=1024, shuffle=False)\n",
        "datapipe = datapipe.sample_neighbor(graph, [4, 4])\n",
        "datapipe = datapipe.fetch_feature(feature, node_feature_keys=[\"feat\"])\n",
        "datapipe = datapipe.to_dgl()\n",
        "datapipe = datapipe.copy_to(device)\n",
        "valid_dataloader = gb.MultiProcessDataLoader(datapipe, num_workers=0)\n",
        "\n",
        "\n",
        "import sklearn.metrics"
      ],
      "metadata": {
        "id": "Gvd7vFWZwQI5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following is a training loop that performs validation every epoch. It also saves the model with the best validation accuracy into a file."
      ],
      "metadata": {
        "id": "nTIIfVMDwXqX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import tqdm\n",
        "\n",
        "best_accuracy = 0\n",
        "best_model_path = \"model.pt\"\n",
        "for epoch in range(10):\n",
        "    model.train()\n",
        "\n",
        "    with tqdm.tqdm(train_dataloader) as tq:\n",
        "        for step, data in enumerate(tq):\n",
        "            x = data.node_features[\"feat\"]\n",
        "            labels = data.labels\n",
        "\n",
        "            predictions = model(data.blocks, x)\n",
        "\n",
        "            loss = F.cross_entropy(predictions, labels)\n",
        "            opt.zero_grad()\n",
        "            loss.backward()\n",
        "            opt.step()\n",
        "\n",
        "            accuracy = sklearn.metrics.accuracy_score(\n",
        "                labels.cpu().numpy(),\n",
        "                predictions.argmax(1).detach().cpu().numpy(),\n",
        "            )\n",
        "\n",
        "            tq.set_postfix(\n",
        "                {\"loss\": \"%.03f\" % loss.item(), \"acc\": \"%.03f\" % accuracy},\n",
        "                refresh=False,\n",
        "            )\n",
        "\n",
        "    model.eval()\n",
        "\n",
        "    predictions = []\n",
        "    labels = []\n",
        "    with tqdm.tqdm(valid_dataloader) as tq, torch.no_grad():\n",
        "        for data in tq:\n",
        "            x = data.node_features[\"feat\"]\n",
        "            labels.append(data.labels.cpu().numpy())\n",
        "            predictions.append(model(data.blocks, x).argmax(1).cpu().numpy())\n",
        "        predictions = np.concatenate(predictions)\n",
        "        labels = np.concatenate(labels)\n",
        "        accuracy = sklearn.metrics.accuracy_score(labels, predictions)\n",
        "        print(\"Epoch {} Validation Accuracy {}\".format(epoch, accuracy))\n",
        "        if best_accuracy < accuracy:\n",
        "            best_accuracy = accuracy\n",
        "            torch.save(model.state_dict(), best_model_path)\n",
        "\n",
        "        # Note that this tutorial do not train the whole model to the end.\n",
        "        break"
      ],
      "metadata": {
        "id": "wsfqhKUvwZEj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Conclusion\n",
        "\n",
        "In this tutorial, you have learned how to train a multi-layer GraphSAGE with neighbor sampling.\n"
      ],
      "metadata": {
        "id": "kmHnUI0QwfJ4"
      }
    }
  ]
}