GNN_Tutorial_NyAI.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/gagan-iitb/CS550/blob/main/GNN_Tutorial_NyAI.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "840e104a",
      "metadata": {
        "id": "840e104a"
      },
      "source": [
        "\n",
        "How Does DGL Represent A Graph?\n",
        "===============================\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0de846d0",
      "metadata": {
        "id": "0de846d0"
      },
      "source": [
        "DGL Graph Construction\n",
        "----------------------\n",
        "\n",
        "DGL represents a directed graph as a ``DGLGraph`` object. You can\n",
        "construct a graph by specifying the number of nodes in the graph as well\n",
        "as the list of source and destination nodes.  Nodes in the graph have\n",
        "consecutive IDs starting from 0.\n",
        "\n",
        "For instance, the following code constructs a directed star graph with 5\n",
        "leaves. The center node's ID is 0. The edges go from the\n",
        "center node to the leaves.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2596e1c8",
      "metadata": {
        "id": "2596e1c8"
      },
      "outputs": [],
      "source": [
        "import dgl\n",
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "g = dgl.graph(([0, 0, 0, 0, 0], [1, 2, 3, 4, 5]), num_nodes=6)\n",
        "# Equivalently, PyTorch LongTensors also work.\n",
        "g = dgl.graph((torch.LongTensor([0, 0, 0, 0, 0]), torch.LongTensor([1, 2, 3, 4, 5])), num_nodes=6)\n",
        "\n",
        "# You can omit the number of nodes argument if you can tell the number of nodes from the edge list alone.\n",
        "g = dgl.graph(([0, 0, 0, 0, 0], [1, 2, 3, 4, 5]))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a49da79d",
      "metadata": {
        "id": "a49da79d"
      },
      "source": [
        "Edges in the graph have consecutive IDs starting from 0, and are\n",
        "in the same order as the list of source and destination nodes during\n",
        "creation.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b0fb8e16",
      "metadata": {
        "id": "b0fb8e16"
      },
      "outputs": [],
      "source": [
        "# Print the source and destination nodes of every edge.\n",
        "print(g.edges())"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "898c2b85",
      "metadata": {
        "id": "898c2b85"
      },
      "source": [
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: ``DGLGraph``'s are always directed to best fit the computation\n",
        "   pattern of graph neural networks, where the messages sent\n",
        "   from one node to the other are often different between both\n",
        "   directions. If you want to handle undirected graphs, you may consider\n",
        "   treating it as a bidirectional graph. See [Graph\n",
        "   Transformations](#Graph-Transformations) for an example of making\n",
        "   a bidirectional graph.\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fcb3686e",
      "metadata": {
        "id": "fcb3686e"
      },
      "source": [
        "Assigning Node and Edge Features to Graph\n",
        "-----------------------------------------\n",
        "\n",
        "Many graph data contain attributes on nodes and edges.\n",
        "Although the types of node and edge attributes can be arbitrary in real\n",
        "world, ``DGLGraph`` only accepts attributes stored in tensors (with\n",
        "numerical contents). Consequently, an attribute of all the nodes or\n",
        "edges must have the same shape. In the context of deep learning, those\n",
        "attributes are often called *features*.\n",
        "\n",
        "You can assign and retrieve node and edge features via ``ndata`` and\n",
        "``edata`` interface.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4c6fc37a",
      "metadata": {
        "id": "4c6fc37a"
      },
      "outputs": [],
      "source": [
        "# Assign a 3-dimensional node feature vector for each node.\n",
        "g.ndata['x'] = torch.randn(6, 3)\n",
        "# Assign a 4-dimensional edge feature vector for each edge.\n",
        "g.edata['a'] = torch.randn(5, 4)\n",
        "# Assign a 5x4 node feature matrix for each node.  Node and edge features in DGL can be multi-dimensional.\n",
        "g.ndata['y'] = torch.randn(6, 5, 4)\n",
        "\n",
        "print(g.edata['a'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "eee5ea61",
      "metadata": {
        "id": "eee5ea61"
      },
      "source": [
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: The vast development of deep learning has provided us many\n",
        "   ways to encode various types of attributes into numerical features.\n",
        "   Here are some general suggestions:\n",
        "\n",
        "   -  For categorical attributes (e.g. gender, occupation), consider\n",
        "      converting them to integers or one-hot encoding.\n",
        "   -  For variable length string contents (e.g. news article, quote),\n",
        "      consider applying a language model.\n",
        "   -  For images, consider applying a vision model such as CNNs.\n",
        "\n",
        "You can find plenty of materials on how to encode such attributes\n",
        "   into a tensor in the [PyTorch Deep Learning\n",
        "   Tutorials](https://pytorch.org/tutorials/)\n",
        "\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4062ccd8",
      "metadata": {
        "id": "4062ccd8"
      },
      "source": [
        "Querying Graph Structures\n",
        "-------------------------\n",
        "\n",
        "``DGLGraph`` object provides various methods to query a graph structure.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f1460185",
      "metadata": {
        "id": "f1460185"
      },
      "outputs": [],
      "source": [
        "print(g.num_nodes())\n",
        "print(g.num_edges())\n",
        "# Out degrees of the center node\n",
        "print(g.out_degrees(0))\n",
        "# In degrees of the center node - note that the graph is directed so the in degree should be 0.\n",
        "print(g.in_degrees(0))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "91adad69",
      "metadata": {
        "id": "91adad69"
      },
      "source": [
        "Graph Transformations\n",
        "---------------------\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e1c8799b",
      "metadata": {
        "id": "e1c8799b"
      },
      "source": [
        "DGL provides many APIs to transform a graph to another such as\n",
        "extracting a subgraph:\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "47bb76ee",
      "metadata": {
        "id": "47bb76ee"
      },
      "outputs": [],
      "source": [
        "# Induce a subgraph from node 0, node 1 and node 3 from the original graph.\n",
        "sg1 = g.subgraph([0, 1, 3])\n",
        "# Induce a subgraph from edge 0, edge 1 and edge 3 from the original graph.\n",
        "sg2 = g.edge_subgraph([0, 1, 3])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4a566d21",
      "metadata": {
        "id": "4a566d21"
      },
      "source": [
        "You can obtain the node/edge mapping from the subgraph to the original\n",
        "graph by looking into the node feature ``dgl.NID`` or edge feature\n",
        "``dgl.EID`` in the new graph.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e22b5e61",
      "metadata": {
        "id": "e22b5e61"
      },
      "outputs": [],
      "source": [
        "# The original IDs of each node in sg1\n",
        "print(sg1.ndata[dgl.NID])\n",
        "# The original IDs of each edge in sg1\n",
        "print(sg1.edata[dgl.EID])\n",
        "# The original IDs of each node in sg2\n",
        "print(sg2.ndata[dgl.NID])\n",
        "# The original IDs of each edge in sg2\n",
        "print(sg2.edata[dgl.EID])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4e173c81",
      "metadata": {
        "id": "4e173c81"
      },
      "source": [
        "``subgraph`` and ``edge_subgraph`` also copies the original features\n",
        "to the subgraph:\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "055798db",
      "metadata": {
        "id": "055798db"
      },
      "outputs": [],
      "source": [
        "# The original node feature of each node in sg1\n",
        "print(sg1.ndata['x'])\n",
        "# The original edge feature of each node in sg1\n",
        "print(sg1.edata['a'])\n",
        "# The original node feature of each node in sg2\n",
        "print(sg2.ndata['x'])\n",
        "# The original edge feature of each node in sg2\n",
        "print(sg2.edata['a'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "de45429f",
      "metadata": {
        "id": "de45429f"
      },
      "source": [
        "Another common transformation is to add a reverse edge for each edge in\n",
        "the original graph with ``dgl.add_reverse_edges``.\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: If you have an undirected graph, it is better to convert it\n",
        "   into a bidirectional graph first via adding reverse edges.\n",
        "\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ab1579be",
      "metadata": {
        "id": "ab1579be"
      },
      "outputs": [],
      "source": [
        "newg = dgl.add_reverse_edges(g)\n",
        "newg.edges()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7c4a1fda",
      "metadata": {
        "id": "7c4a1fda"
      },
      "source": [
        "Loading and Saving Graphs\n",
        "-------------------------\n",
        "\n",
        "You can save a graph or a list of graphs via ``dgl.save_graphs`` and\n",
        "load them back with ``dgl.load_graphs``.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "08fdd1aa",
      "metadata": {
        "id": "08fdd1aa"
      },
      "outputs": [],
      "source": [
        "# Save graphs\n",
        "dgl.save_graphs('graph.dgl', g)\n",
        "dgl.save_graphs('graphs.dgl', [g, sg1, sg2])\n",
        "\n",
        "# Load graphs\n",
        "(g,), _ = dgl.load_graphs('graph.dgl')\n",
        "print(g)\n",
        "(g, sg1, sg2), _ = dgl.load_graphs('graphs.dgl')\n",
        "print(g)\n",
        "print(sg1)\n",
        "print(sg2)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f755ca48",
      "metadata": {
        "id": "f755ca48"
      },
      "source": [
        "\n",
        "Node Classification with DGL\n",
        "============================\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "70fe1b2d",
      "metadata": {
        "id": "70fe1b2d"
      },
      "outputs": [],
      "source": [
        "import dgl\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "949d49ec",
      "metadata": {
        "id": "949d49ec"
      },
      "source": [
        "Overview of Node Classification with GNN\n",
        "----------------------------------------\n",
        "\n",
        "One of the most popular and widely adopted tasks on graph data is node\n",
        "classification, where a model needs to predict the ground truth category\n",
        "of each node. Before graph neural networks, many proposed methods are\n",
        "using either connectivity alone (such as DeepWalk or node2vec), or simple\n",
        "combinations of connectivity and the node's own features.  GNNs, by\n",
        "contrast, offers an opportunity to obtain node representations by\n",
        "combining the connectivity and features of a *local neighborhood*.\n",
        "\n",
        "[Kipf et\n",
        "al.,](https://arxiv.org/abs/1609.02907) is an example that formulates\n",
        "the node classification problem as a semi-supervised node classification\n",
        "task. With the help of only a small portion of labeled nodes, a graph\n",
        "neural network (GNN) can accurately predict the node category of the\n",
        "others.\n",
        "\n",
        "This tutorial will show how to build such a GNN for semi-supervised node\n",
        "classification with only a small number of labels on the **Cora\n",
        "dataset**,\n",
        "a citation network with papers as nodes and citations as edges. The task\n",
        "is to predict the category of a given paper. Each paper node contains a\n",
        "word count vector as its features, normalized so that they sum up to one,\n",
        "as described in Section 5.2 of [the paper](https://arxiv.org/abs/1609.02907).\n",
        "\n",
        "Loading Cora Dataset\n",
        "--------------------\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "093a5e1f",
      "metadata": {
        "id": "093a5e1f"
      },
      "outputs": [],
      "source": [
        "import dgl.data\n",
        "\n",
        "dataset = dgl.data.CoraGraphDataset()\n",
        "print('Number of categories:', dataset.num_classes)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8e6ec6e6",
      "metadata": {
        "id": "8e6ec6e6"
      },
      "source": [
        "A DGL Dataset object may contain one or multiple graphs. The Cora\n",
        "dataset used in this tutorial only consists of one single graph.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7a3c81b6",
      "metadata": {
        "id": "7a3c81b6"
      },
      "outputs": [],
      "source": [
        "g = dataset[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "12b2cbc6",
      "metadata": {
        "id": "12b2cbc6"
      },
      "source": [
        "Here, `g` is a `DGLGraph` object. A `DGLGraph` represents a graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2289c331",
      "metadata": {
        "id": "2289c331"
      },
      "outputs": [],
      "source": [
        "# Get the number of nodes\n",
        "print('Number of nodes:', g.num_nodes())\n",
        "# Get the number of edges\n",
        "print('Number of edges:', g.num_edges())"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7fac54db",
      "metadata": {
        "id": "7fac54db"
      },
      "source": [
        "A `DGLGraph` stores node features and edge features in two\n",
        "dictionary-like attributes called ``ndata`` and ``edata``.\n",
        "In the DGL Cora dataset, the graph contains the following node features:\n",
        "\n",
        "- ``train_mask``: A boolean tensor indicating whether the node is in the\n",
        "  training set.\n",
        "\n",
        "- ``val_mask``: A boolean tensor indicating whether the node is in the\n",
        "  validation set.\n",
        "\n",
        "- ``test_mask``: A boolean tensor indicating whether the node is in the\n",
        "  test set.\n",
        "\n",
        "- ``label``: The ground truth node category.\n",
        "\n",
        "-  ``feat``: The node features.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "41fe5482",
      "metadata": {
        "id": "41fe5482"
      },
      "outputs": [],
      "source": [
        "print('Node feature names:', g.ndata.keys())\n",
        "print('Edge feature names:', g.edata.keys())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ca97e54f",
      "metadata": {
        "id": "ca97e54f"
      },
      "outputs": [],
      "source": [
        "g.ndata['train_mask']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2f85fef3",
      "metadata": {
        "id": "2f85fef3"
      },
      "outputs": [],
      "source": [
        "print('Number of training nodes:', int(g.ndata['train_mask'].int().sum()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "669822b2",
      "metadata": {
        "id": "669822b2"
      },
      "outputs": [],
      "source": [
        "g.ndata['label']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d209c39",
      "metadata": {
        "id": "9d209c39"
      },
      "outputs": [],
      "source": [
        "print('Node feature tensor shape:', g.ndata['feat'].shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9a4ad944",
      "metadata": {
        "id": "9a4ad944"
      },
      "source": [
        "Defining a Graph Convolutional Network (GCN)\n",
        "--------------------------------------------\n",
        "\n",
        "This tutorial will build a two-layer [Graph Convolutional Network(GCN)](http://tkipf.github.io/graph-convolutional-networks/). Each\n",
        "layer computes new node representations by aggregating neighbor\n",
        "information.\n",
        "\n",
        "![img](https://tkipf.github.io/graph-convolutional-networks/images/gcn_web.png)\n",
        "\n",
        "To build a multi-layer GCN you can simply stack ``dgl.nn.GraphConv``\n",
        "modules, which inherit ``torch.nn.Module``.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a50b99bc",
      "metadata": {
        "id": "a50b99bc"
      },
      "outputs": [],
      "source": [
        "from dgl.nn import GraphConv\n",
        "\n",
        "class GCN(nn.Module):\n",
        "    def __init__(self, in_feats, h_feats, num_classes):\n",
        "        super(GCN, self).__init__()\n",
        "        self.conv1 = GraphConv(in_feats, h_feats)\n",
        "        self.conv2 = GraphConv(h_feats, num_classes)\n",
        "\n",
        "    def forward(self, g, in_feat):\n",
        "        h = self.conv1(g, in_feat)\n",
        "        h = F.relu(h)\n",
        "        h = self.conv2(g, h)\n",
        "        return h\n",
        "\n",
        "# Create the model with given dimensions\n",
        "model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f1ac1029",
      "metadata": {
        "id": "f1ac1029"
      },
      "source": [
        "DGL provides implementation of many popular neighbor aggregation\n",
        "modules. You can easily invoke them with one line of code.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1375f81a",
      "metadata": {
        "id": "1375f81a"
      },
      "source": [
        "Training the GCN\n",
        "----------------\n",
        "\n",
        "Training this GCN is similar to training other PyTorch neural networks.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7564eeae",
      "metadata": {
        "id": "7564eeae"
      },
      "outputs": [],
      "source": [
        "def train(g, model):\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
        "    best_val_acc = 0\n",
        "    best_test_acc = 0\n",
        "\n",
        "    features = g.ndata['feat']\n",
        "    labels = g.ndata['label']\n",
        "    train_mask = g.ndata['train_mask']\n",
        "    val_mask = g.ndata['val_mask']\n",
        "    test_mask = g.ndata['test_mask']\n",
        "\n",
        "    for e in range(100):\n",
        "        # Forward\n",
        "        logits = model(g, features)\n",
        "\n",
        "        # Compute prediction\n",
        "        pred = logits.argmax(1)\n",
        "\n",
        "        # Compute loss\n",
        "        # Note that you should only compute the losses of the nodes in the training set.\n",
        "        loss = F.cross_entropy(logits[train_mask], labels[train_mask])\n",
        "\n",
        "        # Compute accuracy on training/validation/test\n",
        "        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()\n",
        "        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()\n",
        "        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()\n",
        "\n",
        "        # Save the best validation accuracy and the corresponding test accuracy.\n",
        "        if best_val_acc < val_acc:\n",
        "            best_val_acc = val_acc\n",
        "            best_test_acc = test_acc\n",
        "\n",
        "        # Backward\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if e % 5 == 0:\n",
        "            print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(\n",
        "                e, loss, val_acc, best_val_acc, test_acc, best_test_acc))\n",
        "model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes)\n",
        "train(g, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bf208e05",
      "metadata": {
        "id": "bf208e05"
      },
      "source": [
        "Training on GPU\n",
        "---------------\n",
        "\n",
        "Training on GPU requires to put both the model and the graph onto GPU\n",
        "with the ``to`` method, similar to what you will do in PyTorch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5f9c45a9",
      "metadata": {
        "id": "5f9c45a9"
      },
      "outputs": [],
      "source": [
        "print('Before:', g.device)\n",
        "g = g.to('cuda')\n",
        "print('After:', g.device)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d276c794",
      "metadata": {
        "id": "d276c794"
      },
      "source": [
        "It copies all the `ndata` and `edata` to GPU too."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7b8c4078",
      "metadata": {
        "id": "7b8c4078"
      },
      "outputs": [],
      "source": [
        "print(g.ndata['train_mask'].device)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d435287d",
      "metadata": {
        "id": "d435287d"
      },
      "source": [
        "Train it on GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e4c7fe5e",
      "metadata": {
        "id": "e4c7fe5e"
      },
      "outputs": [],
      "source": [
        "model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')\n",
        "train(g, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dbbfba7b",
      "metadata": {
        "id": "dbbfba7b"
      },
      "source": [
        "\n",
        "Write your own GNN module\n",
        "========================="
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f60a9a80",
      "metadata": {
        "id": "f60a9a80"
      },
      "outputs": [],
      "source": [
        "import dgl\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "81d5f2b6",
      "metadata": {
        "id": "81d5f2b6"
      },
      "source": [
        "Message passing and GNNs\n",
        "------------------------\n",
        "\n",
        "DGL follows the *message passing paradigm* inspired by the Message\n",
        "Passing Neural Network proposed by [Gilmer et\n",
        "al.](https://arxiv.org/abs/1704.01212) Essentially, they found many\n",
        "GNN models can fit into the following framework:\n",
        "\n",
        "$$m_{u\\to v}^{(l)} = M^{(l)}\\left(h_v^{(l-1)}, h_u^{(l-1)}, e_{u\\to v}^{(l-1)}\\right)$$\n",
        "\n",
        "\n",
        "$$m_{v}^{(l)} = \\sum_{u\\in\\mathcal{N}(v)}m_{u\\to v}^{(l)};$$\n",
        "\n",
        "$$h_v^{(l)} = U^{(l)}\\left(h_v^{(l-1)}, m_v^{(l)}\\right)\n",
        "$$\n",
        "\n",
        "where DGL calls $M^{(l)}$ the *message function*, $\\sum$ the\n",
        "*reduce function* and $U^{(l)}$ the *update function*. Note that\n",
        "$\\sum$ here can represent any function and is not necessarily a\n",
        "summation.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a5e6f0a8",
      "metadata": {
        "id": "a5e6f0a8"
      },
      "source": [
        "For example, the [GraphSAGE convolution (Hamilton et al.,\n",
        "2017)](https://cs.stanford.edu/people/jure/pubs/graphsage-nips17.pdf)\n",
        "takes the following mathematical form:\n",
        "\n",
        "$$\n",
        "h_{\\mathcal{N}(v)}^k\\leftarrow \\text{Average}\\{h_u^{k-1},\\forall u\\in\\mathcal{N}(v)\\}\\\\\n",
        "h_v^k\\leftarrow \\text{ReLU}\\left(W^k\\cdot \\text{CONCAT}(h_v^{k-1}, h_{\\mathcal{N}(v)}^k) \\right)\n",
        "$$\n",
        "\n",
        "You can see that message passing is directional: the message sent from\n",
        "one node $u$ to other node $v$ is not necessarily the same\n",
        "as the other message sent from node $v$ to node $u$ in the\n",
        "opposite direction.\n",
        "\n",
        "Although DGL has builtin support of GraphSAGE via `dgl.nn.pytorch.SAGEConv`,\n",
        "here is how you can implement GraphSAGE convolution in DGL by your own."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ebcc8ed5",
      "metadata": {
        "id": "ebcc8ed5"
      },
      "outputs": [],
      "source": [
        "import dgl.function as fn\n",
        "\n",
        "class SAGEConv(nn.Module):\n",
        "    \"\"\"Graph convolution module used by the GraphSAGE model.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    in_feat : int\n",
        "        Input feature size.\n",
        "    out_feat : int\n",
        "        Output feature size.\n",
        "    \"\"\"\n",
        "    def __init__(self, in_feat, out_feat):\n",
        "        super(SAGEConv, self).__init__()\n",
        "        # A linear submodule for projecting the input and neighbor feature to the output.\n",
        "        self.linear = nn.Linear(in_feat * 2, out_feat)\n",
        "\n",
        "    def forward(self, g, h):\n",
        "        \"\"\"Forward computation\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        g : Graph\n",
        "            The input graph.\n",
        "        h : Tensor\n",
        "            The input node feature.\n",
        "        \"\"\"\n",
        "        with g.local_scope():\n",
        "            g.ndata['h'] = h\n",
        "            # update_all is a message passing API.\n",
        "            g.update_all(message_func=fn.copy_u('h', 'm'), reduce_func=fn.mean('m', 'h_N'))\n",
        "            h_N = g.ndata['h_N']\n",
        "            h_total = torch.cat([h, h_N], dim=1)\n",
        "            return self.linear(h_total)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a8f2aa34",
      "metadata": {
        "id": "a8f2aa34"
      },
      "source": [
        "The central piece in this code is the `g.update_all`\n",
        "function, which gathers and averages the neighbor features. There are\n",
        "three concepts here:\n",
        "\n",
        "* Message function ``fn.copy_u('h', 'm')`` that\n",
        "  copies the node feature under name ``'h'`` as *messages* sent to\n",
        "  neighbors.\n",
        "\n",
        "* Reduce function ``fn.mean('m', 'h_N')`` that averages\n",
        "  all the received messages under name ``'m'`` and saves the result as a\n",
        "  new node feature ``'h_N'``.\n",
        "\n",
        "* ``update_all`` tells DGL to trigger the\n",
        "  message and reduce functions for all the nodes and edges.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d57d6ab7",
      "metadata": {
        "id": "d57d6ab7"
      },
      "source": [
        "Afterwards, you can stack your own GraphSAGE convolution layers to form\n",
        "a multi-layer GraphSAGE network.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8e0e7815",
      "metadata": {
        "id": "8e0e7815"
      },
      "outputs": [],
      "source": [
        "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)\n",
        "        self.conv2 = SAGEConv(h_feats, num_classes)\n",
        "\n",
        "    def forward(self, g, in_feat):\n",
        "        h = self.conv1(g, in_feat)\n",
        "        h = F.relu(h)\n",
        "        h = self.conv2(g, h)\n",
        "        return h"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e75e9135",
      "metadata": {
        "id": "e75e9135"
      },
      "source": [
        "### Training loop\n",
        "\n",
        "The following code for data loading and training loop is directly copied\n",
        "from the introduction tutorial.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "24b6bc77",
      "metadata": {
        "id": "24b6bc77"
      },
      "outputs": [],
      "source": [
        "import dgl.data\n",
        "\n",
        "dataset = dgl.data.CoraGraphDataset()\n",
        "g = dataset[0]\n",
        "\n",
        "def train(g, model):\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
        "    all_logits = []\n",
        "    best_val_acc = 0\n",
        "    best_test_acc = 0\n",
        "\n",
        "    features = g.ndata['feat']\n",
        "    labels = g.ndata['label']\n",
        "    train_mask = g.ndata['train_mask']\n",
        "    val_mask = g.ndata['val_mask']\n",
        "    test_mask = g.ndata['test_mask']\n",
        "    for e in range(200):\n",
        "        # Forward\n",
        "        logits = model(g, features)\n",
        "\n",
        "        # Compute prediction\n",
        "        pred = logits.argmax(1)\n",
        "\n",
        "        # Compute loss\n",
        "        # Note that we should only compute the losses of the nodes in the training set,\n",
        "        # i.e. with train_mask 1.\n",
        "        loss = F.cross_entropy(logits[train_mask], labels[train_mask])\n",
        "\n",
        "        # Compute accuracy on training/validation/test\n",
        "        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()\n",
        "        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()\n",
        "        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()\n",
        "\n",
        "        # Save the best validation accuracy and the corresponding test accuracy.\n",
        "        if best_val_acc < val_acc:\n",
        "            best_val_acc = val_acc\n",
        "            best_test_acc = test_acc\n",
        "\n",
        "        # Backward\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        all_logits.append(logits.detach())\n",
        "\n",
        "        if e % 5 == 0:\n",
        "            print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(\n",
        "                e, loss, val_acc, best_val_acc, test_acc, best_test_acc))\n",
        "\n",
        "model = Model(g.ndata['feat'].shape[1], 16, dataset.num_classes)\n",
        "train(g, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a5f8dff3",
      "metadata": {
        "id": "a5f8dff3"
      },
      "source": [
        "More customization\n",
        "------------------\n",
        "\n",
        "In DGL, we provide many built-in message and reduce functions under the\n",
        "``dgl.function`` package. You can find more details in [the API\n",
        "documentation](https://docs.dgl.ai/api/python/dgl.function.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9b3d933c",
      "metadata": {
        "id": "9b3d933c"
      },
      "source": [
        "These APIs allow one to quickly implement new graph convolution modules.\n",
        "For example, the following implements a new ``SAGEConv`` that aggregates\n",
        "neighbor representations using a weighted average. Note that ``edata``\n",
        "member can hold edge features which can also take part in message\n",
        "passing.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8ec5cdcd",
      "metadata": {
        "id": "8ec5cdcd"
      },
      "outputs": [],
      "source": [
        "class WeightedSAGEConv(nn.Module):\n",
        "    \"\"\"Graph convolution module used by the GraphSAGE model with edge weights.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    in_feat : int\n",
        "        Input feature size.\n",
        "    out_feat : int\n",
        "        Output feature size.\n",
        "    \"\"\"\n",
        "    def __init__(self, in_feat, out_feat):\n",
        "        super(WeightedSAGEConv, self).__init__()\n",
        "        # A linear submodule for projecting the input and neighbor feature to the output.\n",
        "        self.linear = nn.Linear(in_feat * 2, out_feat)\n",
        "\n",
        "    def forward(self, g, h, w):\n",
        "        \"\"\"Forward computation\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        g : Graph\n",
        "            The input graph.\n",
        "        h : Tensor\n",
        "            The input node feature.\n",
        "        w : Tensor\n",
        "            The edge weight.\n",
        "        \"\"\"\n",
        "        with g.local_scope():\n",
        "            g.ndata['h'] = h\n",
        "            g.edata['w'] = w\n",
        "            g.update_all(message_func=fn.u_mul_e('h', 'w', 'm'), reduce_func=fn.mean('m', 'h_N'))\n",
        "            h_N = g.ndata['h_N']\n",
        "            h_total = torch.cat([h, h_N], dim=1)\n",
        "            return self.linear(h_total)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "155f7ab0",
      "metadata": {
        "id": "155f7ab0"
      },
      "source": [
        "Because the graph in this dataset does not have edge weights, we\n",
        "manually assign all edge weights to one in the ``forward()`` function of\n",
        "the model. You can replace it with your own edge weights.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "23c04135",
      "metadata": {
        "id": "23c04135"
      },
      "outputs": [],
      "source": [
        "class Model(nn.Module):\n",
        "    def __init__(self, in_feats, h_feats, num_classes):\n",
        "        super(Model, self).__init__()\n",
        "        self.conv1 = WeightedSAGEConv(in_feats, h_feats)\n",
        "        self.conv2 = WeightedSAGEConv(h_feats, num_classes)\n",
        "\n",
        "    def forward(self, g, in_feat):\n",
        "        h = self.conv1(g, in_feat, torch.ones(g.num_edges()).to(g.device))\n",
        "        h = F.relu(h)\n",
        "        h = self.conv2(g, h, torch.ones(g.num_edges()).to(g.device))\n",
        "        return h\n",
        "\n",
        "model = Model(g.ndata['feat'].shape[1], 16, dataset.num_classes)\n",
        "train(g, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "739b7700",
      "metadata": {
        "id": "739b7700"
      },
      "source": [
        "Even more customization by user-defined function\n",
        "------------------------------------------------\n",
        "\n",
        "DGL allows user-defined message and reduce function for the maximal\n",
        "expressiveness. Here is a user-defined message function that is\n",
        "equivalent to ``fn.u_mul_e('h', 'w', 'm')``.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4424f8d0",
      "metadata": {
        "id": "4424f8d0"
      },
      "outputs": [],
      "source": [
        "def u_mul_e_udf(edges):\n",
        "    return {'m' : edges.src['h'] * edges.data['w']}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ffcff042",
      "metadata": {
        "id": "ffcff042"
      },
      "source": [
        "``edges`` has three members: ``src``, ``data`` and ``dst``, representing\n",
        "the source node feature, edge feature, and destination node feature for\n",
        "all edges.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a7bec7f1",
      "metadata": {
        "id": "a7bec7f1"
      },
      "source": [
        "You can also write your own reduce function. For example, the following\n",
        "is equivalent to the builtin ``fn.sum('m', 'h')`` function that sums up\n",
        "the incoming messages:\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d71a9cfa",
      "metadata": {
        "id": "d71a9cfa"
      },
      "outputs": [],
      "source": [
        "def sum_udf(nodes):\n",
        "    return {'h': nodes.mailbox['m'].sum(1)}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5f935673",
      "metadata": {
        "id": "5f935673"
      },
      "source": [
        "In short, DGL will group the nodes by their in-degrees, and for each\n",
        "group DGL stacks the incoming messages along the second dimension. You\n",
        "can then perform a reduction along the second dimension to aggregate\n",
        "messages.\n",
        "\n",
        "For more details on customizing message and reduce function with\n",
        "user-defined function, please refer to the [API\n",
        "reference](https://docs.dgl.ai/api/python/udf.html).\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "399be5a3",
      "metadata": {
        "id": "399be5a3"
      },
      "source": [
        "Best practice of writing custom GNN modules\n",
        "-------------------------------------------\n",
        "\n",
        "DGL recommends the following practice ranked by preference:\n",
        "\n",
        "-  Use ``dgl.nn`` modules.\n",
        "-  Use ``dgl.nn.functional`` functions which contain lower-level complex\n",
        "   operations such as computing a softmax for each node over incoming\n",
        "   edges.\n",
        "-  Use ``update_all`` with builtin message and reduce functions.\n",
        "-  Use user-defined message or reduce functions.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "efceacd9",
      "metadata": {
        "id": "efceacd9"
      },
      "source": [
        "\n",
        "Link Prediction using Graph Neural Networks\n",
        "==========================================="
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f1a9f52d",
      "metadata": {
        "id": "f1a9f52d"
      },
      "outputs": [],
      "source": [
        "import dgl\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import itertools\n",
        "import numpy as np\n",
        "import scipy.sparse as sp"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "accef7d8",
      "metadata": {
        "id": "accef7d8"
      },
      "source": [
        "Overview of Link Prediction with GNN\n",
        "------------------------------------\n",
        "\n",
        "Many applications such as social recommendation, item recommendation,\n",
        "knowledge graph completion, etc., can be formulated as link prediction,\n",
        "which predicts whether an edge exists between two particular nodes. This\n",
        "tutorial shows an example of predicting whether a citation relationship,\n",
        "either citing or being cited, between two papers exists in a citation\n",
        "network.\n",
        "\n",
        "This tutorial formulates the link prediction problem as a binary classification\n",
        "problem as follows:\n",
        "\n",
        "-  Treat the edges in the graph as *positive examples*.\n",
        "-  Sample a number of non-existent edges (i.e. node pairs with no edges\n",
        "   between them) as *negative* examples.\n",
        "-  Divide the positive examples and negative examples into a training\n",
        "   set and a test set.\n",
        "-  Evaluate the model with any binary classification metric such as Area\n",
        "   Under Curve (AUC).\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: The practice comes from\n",
        "   [SEAL](https://papers.nips.cc/paper/2018/file/53f0d7c537d99b3824f0f99d62ea2428-Paper.pdf),\n",
        "   although the model here does not use their idea of node labeling.\n",
        "\n",
        "</div>\n",
        "\n",
        "In some domains such as large-scale recommender systems or information\n",
        "retrieval, you may favor metrics that emphasize good performance of\n",
        "top-K predictions. In these cases you may want to consider other metrics\n",
        "such as mean average precision, and use other negative sampling methods,\n",
        "which are beyond the scope of this tutorial.\n",
        "\n",
        "Loading graph and features\n",
        "--------------------------\n",
        "\n",
        "Following the [introduction](1_introduction.ipynb), this tutorial\n",
        "first loads the Cora dataset.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4f8af061",
      "metadata": {
        "id": "4f8af061"
      },
      "outputs": [],
      "source": [
        "import dgl.data\n",
        "\n",
        "dataset = dgl.data.CoraGraphDataset()\n",
        "g = dataset[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7634ed02",
      "metadata": {
        "id": "7634ed02"
      },
      "source": [
        "Prepare training and testing sets\n",
        "---------------------------------\n",
        "\n",
        "This tutorial randomly picks 10% of the edges for positive examples in\n",
        "the test set, and leave the rest for the training set. It then samples\n",
        "the same number of edges for negative examples in both sets.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee117dd3",
      "metadata": {
        "id": "ee117dd3"
      },
      "outputs": [],
      "source": [
        "# Split edge set for training and testing\n",
        "u, v = g.edges()\n",
        "\n",
        "eids = np.arange(g.number_of_edges())\n",
        "eids = np.random.permutation(eids)\n",
        "test_size = int(len(eids) * 0.1)\n",
        "train_size = g.number_of_edges() - test_size\n",
        "test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]\n",
        "train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]\n",
        "\n",
        "# Find all negative edges and split them for training and testing\n",
        "adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy())))\n",
        "adj_neg = 1 - adj.todense() - np.eye(g.number_of_nodes())\n",
        "neg_u, neg_v = np.where(adj_neg != 0)\n",
        "\n",
        "neg_eids = np.random.choice(len(neg_u), g.number_of_edges() // 2)\n",
        "test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]\n",
        "train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e5c00895",
      "metadata": {
        "id": "e5c00895"
      },
      "source": [
        "When training, you will need to remove the edges in the test set from\n",
        "the original graph. You can do this via ``dgl.remove_edges``.\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: ``dgl.remove_edges`` works by creating a subgraph from the\n",
        "   original graph, resulting in a copy and therefore could be slow for\n",
        "   large graphs. If so, you could save the training and test graph to\n",
        "   disk, as you would do for preprocessing.\n",
        "\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c49a25eb",
      "metadata": {
        "id": "c49a25eb"
      },
      "outputs": [],
      "source": [
        "train_g = dgl.remove_edges(g, eids[:test_size])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b29bf1df",
      "metadata": {
        "id": "b29bf1df"
      },
      "source": [
        "Define a GraphSAGE model\n",
        "------------------------\n",
        "\n",
        "This tutorial builds a model consisting of two\n",
        "[GraphSAGE](https://arxiv.org/abs/1706.02216) layers, each computes\n",
        "new node representations by averaging neighbor information. DGL provides\n",
        "``dgl.nn.SAGEConv`` that conveniently creates a GraphSAGE layer.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c34f2fb4",
      "metadata": {
        "id": "c34f2fb4"
      },
      "outputs": [],
      "source": [
        "from dgl.nn import SAGEConv\n",
        "\n",
        "# ----------- 2. create model -------------- #\n",
        "# build a two-layer GraphSAGE model\n",
        "class GraphSAGE(nn.Module):\n",
        "    def __init__(self, in_feats, h_feats):\n",
        "        super(GraphSAGE, self).__init__()\n",
        "        self.conv1 = SAGEConv(in_feats, h_feats, 'mean')\n",
        "        self.conv2 = SAGEConv(h_feats, h_feats, 'mean')\n",
        "\n",
        "    def forward(self, g, in_feat):\n",
        "        h = self.conv1(g, in_feat)\n",
        "        h = F.relu(h)\n",
        "        h = self.conv2(g, h)\n",
        "        return h"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "eae0ca13",
      "metadata": {
        "id": "eae0ca13"
      },
      "source": [
        "The model then predicts the probability of existence of an edge by\n",
        "computing a score between the representations of both incident nodes\n",
        "with a function (e.g. an MLP or a dot product), which you will see in\n",
        "the next section.\n",
        "\n",
        "\\begin{align}\\hat{y}_{u\\sim v} = f(h_u, h_v)\\end{align}\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e64beabb",
      "metadata": {
        "id": "e64beabb"
      },
      "source": [
        "Positive graph, negative graph, and ``apply_edges``\n",
        "---------------------------------------------------\n",
        "\n",
        "In previous tutorials you have learned how to compute node\n",
        "representations with a GNN. However, link prediction requires you to\n",
        "compute representation of *pairs of nodes*.\n",
        "\n",
        "DGL recommends you to treat the pairs of nodes as another graph, since\n",
        "you can describe a pair of nodes with an edge. In link prediction, you\n",
        "will have a *positive graph* consisting of all the positive examples as\n",
        "edges, and a *negative graph* consisting of all the negative examples.\n",
        "The *positive graph* and the *negative graph* will contain the same set\n",
        "of nodes as the original graph.  This makes it easier to pass node\n",
        "features among multiple graphs for computation.  As you will see later,\n",
        "you can directly fed the node representations computed on the entire\n",
        "graph to the positive and the negative graphs for computing pair-wise\n",
        "scores.\n",
        "\n",
        "The following code constructs the positive graph and the negative graph\n",
        "for the training set and the test set respectively.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d2a964b",
      "metadata": {
        "id": "9d2a964b"
      },
      "outputs": [],
      "source": [
        "train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.number_of_nodes())\n",
        "train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.number_of_nodes())\n",
        "\n",
        "test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.number_of_nodes())\n",
        "test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.number_of_nodes())"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ee8975ad",
      "metadata": {
        "id": "ee8975ad"
      },
      "source": [
        "The benefit of treating the pairs of nodes as a graph is that you can\n",
        "use the ``DGLGraph.apply_edges`` method, which conveniently computes new\n",
        "edge features based on the incident nodes’ features and the original\n",
        "edge features (if applicable).\n",
        "\n",
        "DGL provides a set of optimized builtin functions to compute new\n",
        "edge features based on the original node/edge features. For example,\n",
        "``dgl.function.u_dot_v`` computes a dot product of the incident nodes’\n",
        "representations for each edge.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d850a7bc",
      "metadata": {
        "id": "d850a7bc"
      },
      "outputs": [],
      "source": [
        "import dgl.function as fn\n",
        "\n",
        "class DotPredictor(nn.Module):\n",
        "    def forward(self, g, h):\n",
        "        with g.local_scope():\n",
        "            g.ndata['h'] = h\n",
        "            # Compute a new edge feature named 'score' by a dot-product between the\n",
        "            # source node feature 'h' and destination node feature 'h'.\n",
        "            g.apply_edges(fn.u_dot_v('h', 'h', 'score'))\n",
        "            # u_dot_v returns a 1-element vector for each edge so you need to squeeze it.\n",
        "            return g.edata['score'][:, 0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "402a1512",
      "metadata": {
        "id": "402a1512"
      },
      "source": [
        "You can also write your own function if it is complex.\n",
        "For instance, the following module produces a scalar score on each edge\n",
        "by concatenating the incident nodes’ features and passing it to an MLP.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bd0fef6d",
      "metadata": {
        "id": "bd0fef6d"
      },
      "outputs": [],
      "source": [
        "class MLPPredictor(nn.Module):\n",
        "    def __init__(self, h_feats):\n",
        "        super().__init__()\n",
        "        self.W1 = nn.Linear(h_feats * 2, h_feats)\n",
        "        self.W2 = nn.Linear(h_feats, 1)\n",
        "\n",
        "    def apply_edges(self, edges):\n",
        "        \"\"\"\n",
        "        Computes a scalar score for each edge of the given graph.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        edges :\n",
        "            Has three members ``src``, ``dst`` and ``data``, each of\n",
        "            which is a dictionary representing the features of the\n",
        "            source nodes, the destination nodes, and the edges\n",
        "            themselves.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        dict\n",
        "            A dictionary of new edge features.\n",
        "        \"\"\"\n",
        "        h = torch.cat([edges.src['h'], edges.dst['h']], 1)\n",
        "        return {'score': self.W2(F.relu(self.W1(h))).squeeze(1)}\n",
        "\n",
        "    def forward(self, g, h):\n",
        "        with g.local_scope():\n",
        "            g.ndata['h'] = h\n",
        "            g.apply_edges(self.apply_edges)\n",
        "            return g.edata['score']"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e470aea8",
      "metadata": {
        "id": "e470aea8"
      },
      "source": [
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: The builtin functions are optimized for both speed and memory.\n",
        "   We recommend using builtin functions whenever possible.\n",
        "\n",
        "</div>\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: If you have read the [message passing\n",
        "   tutorial](3_message_passing.ipynb), you will notice that the\n",
        "   argument ``apply_edges`` takes has exactly the same form as a message\n",
        "   function in ``update_all``.\n",
        "\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "65b40145",
      "metadata": {
        "id": "65b40145"
      },
      "source": [
        "Training loop\n",
        "-------------\n",
        "\n",
        "After you defined the node representation computation and the edge score\n",
        "computation, you can go ahead and define the overall model, loss\n",
        "function, and evaluation metric.\n",
        "\n",
        "The loss function is simply binary cross entropy loss.\n",
        "\n",
        "\\begin{align}\\mathcal{L} = -\\sum_{u\\sim v\\in \\mathcal{D}}\\left( y_{u\\sim v}\\log(\\hat{y}_{u\\sim v}) + (1-y_{u\\sim v})\\log(1-\\hat{y}_{u\\sim v})) \\right)\\end{align}\n",
        "\n",
        "The evaluation metric in this tutorial is AUC.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "42e96856",
      "metadata": {
        "id": "42e96856"
      },
      "outputs": [],
      "source": [
        "model = GraphSAGE(train_g.ndata['feat'].shape[1], 16)\n",
        "# You can replace DotPredictor with MLPPredictor.\n",
        "#pred = MLPPredictor(16)\n",
        "pred = DotPredictor()\n",
        "\n",
        "def compute_loss(pos_score, neg_score):\n",
        "    scores = torch.cat([pos_score, neg_score])\n",
        "    labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])])\n",
        "    return F.binary_cross_entropy_with_logits(scores, labels)\n",
        "\n",
        "def compute_auc(pos_score, neg_score):\n",
        "    scores = torch.cat([pos_score, neg_score]).numpy()\n",
        "    labels = torch.cat(\n",
        "        [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()\n",
        "    return roc_auc_score(labels, scores)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d0d028e2",
      "metadata": {
        "id": "d0d028e2"
      },
      "source": [
        "The training loop goes as follows:\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: This tutorial does not include evaluation on a validation\n",
        "   set. In practice you should save and evaluate the best model based on\n",
        "   performance on the validation set.\n",
        "\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8f9fdf3c",
      "metadata": {
        "id": "8f9fdf3c"
      },
      "outputs": [],
      "source": [
        "# ----------- 3. set up loss and optimizer -------------- #\n",
        "# in this case, loss will in training loop\n",
        "optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=0.01)\n",
        "\n",
        "# ----------- 4. training -------------------------------- #\n",
        "all_logits = []\n",
        "for e in range(100):\n",
        "    # forward\n",
        "    h = model(train_g, train_g.ndata['feat'])\n",
        "    pos_score = pred(train_pos_g, h)\n",
        "    neg_score = pred(train_neg_g, h)\n",
        "    loss = compute_loss(pos_score, neg_score)\n",
        "\n",
        "    # backward\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if e % 5 == 0:\n",
        "        print('In epoch {}, loss: {}'.format(e, loss))\n",
        "\n",
        "# ----------- 5. check results ------------------------ #\n",
        "from sklearn.metrics import roc_auc_score\n",
        "with torch.no_grad():\n",
        "    pos_score = pred(test_pos_g, h)\n",
        "    neg_score = pred(test_neg_g, h)\n",
        "    print('AUC', compute_auc(pos_score, neg_score))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "891c05e9",
      "metadata": {
        "id": "891c05e9"
      },
      "source": [
        "\n",
        "Training a GNN for Graph Classification\n",
        "=======================================\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "18786c11",
      "metadata": {
        "id": "18786c11"
      },
      "outputs": [],
      "source": [
        "import dgl\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c517fda2",
      "metadata": {
        "id": "c517fda2"
      },
      "source": [
        "Overview of Graph Classification with GNN\n",
        "-----------------------------------------\n",
        "\n",
        "Graph classification or regression requires a model to predict certain\n",
        "graph-level properties of a single graph given its node and edge\n",
        "features.  Molecular property prediction is one particular application.\n",
        "\n",
        "This tutorial shows how to train a graph classification model for a\n",
        "small dataset from the paper [How Powerful Are Graph Neural\n",
        "Networks](https://arxiv.org/abs/1810.00826).\n",
        "\n",
        "Loading Data\n",
        "------------\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6ced3175",
      "metadata": {
        "id": "6ced3175"
      },
      "outputs": [],
      "source": [
        "import dgl.data\n",
        "\n",
        "# Generate a synthetic dataset with 10000 graphs, ranging from 10 to 500 nodes.\n",
        "dataset = dgl.data.GINDataset('PROTEINS', self_loop=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0afb934a",
      "metadata": {
        "id": "0afb934a"
      },
      "source": [
        "The dataset is a set of graphs, each with node features and a single\n",
        "label. One can see the node feature dimensionality and the number of\n",
        "possible graph categories of ``GINDataset`` objects in ``dim_nfeats``\n",
        "and ``gclasses`` attributes.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "323e2e4e",
      "metadata": {
        "id": "323e2e4e"
      },
      "outputs": [],
      "source": [
        "print('Node feature dimensionality:', dataset.dim_nfeats)\n",
        "print('Number of graph categories:', dataset.gclasses)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ec0aa99d",
      "metadata": {
        "id": "ec0aa99d"
      },
      "source": [
        "Defining Data Loader\n",
        "--------------------\n",
        "\n",
        "A graph classification dataset usually contains two types of elements: a\n",
        "set of graphs, and their graph-level labels. Similar to an image\n",
        "classification task, when the dataset is large enough, we need to train\n",
        "with mini-batches. When you train a model for image classification or\n",
        "language modeling, you will use a ``DataLoader`` to iterate over the\n",
        "dataset. In DGL, you can use the ``GraphDataLoader``.\n",
        "\n",
        "You can also use various dataset samplers provided in\n",
        "[`torch.utils.data.sampler`](https://pytorch.org/docs/stable/data.html#data-loading-order-and-sampler).\n",
        "For example, this tutorial creates a training ``GraphDataLoader`` and\n",
        "test ``GraphDataLoader``, using ``SubsetRandomSampler`` to tell PyTorch\n",
        "to sample from only a subset of the dataset.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "19209cd4",
      "metadata": {
        "id": "19209cd4"
      },
      "outputs": [],
      "source": [
        "from dgl.dataloading import GraphDataLoader\n",
        "from torch.utils.data.sampler import SubsetRandomSampler\n",
        "\n",
        "num_examples = len(dataset)\n",
        "num_train = int(num_examples * 0.8)\n",
        "\n",
        "train_sampler = SubsetRandomSampler(torch.arange(num_train))\n",
        "test_sampler = SubsetRandomSampler(torch.arange(num_train, num_examples))\n",
        "\n",
        "train_dataloader = GraphDataLoader(\n",
        "    dataset, sampler=train_sampler, batch_size=5, drop_last=False)\n",
        "test_dataloader = GraphDataLoader(\n",
        "    dataset, sampler=test_sampler, batch_size=5, drop_last=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "647d3324",
      "metadata": {
        "id": "647d3324"
      },
      "source": [
        "You can try to iterate over the created ``GraphDataLoader`` and see what it\n",
        "gives:\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6aae5360",
      "metadata": {
        "id": "6aae5360"
      },
      "outputs": [],
      "source": [
        "it = iter(train_dataloader)\n",
        "batch = next(it)\n",
        "print(batch)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d12c49ac",
      "metadata": {
        "id": "d12c49ac"
      },
      "source": [
        "As each element in ``dataset`` has a graph and a label, the\n",
        "``GraphDataLoader`` will return two objects for each iteration. The\n",
        "first element is the batched graph, and the second element is simply a\n",
        "label vector representing the category of each graph in the mini-batch.\n",
        "Next, we’ll talked about the batched graph.\n",
        "\n",
        "A Batched Graph in DGL\n",
        "----------------------\n",
        "\n",
        "In each mini-batch, the sampled graphs are combined into a single bigger\n",
        "batched graph via ``dgl.batch``. The single bigger batched graph merges\n",
        "all original graphs as separately connected components, with the node\n",
        "and edge features concatenated. This bigger graph is also a ``DGLGraph``\n",
        "instance (so you can\n",
        "still treat it as a normal ``DGLGraph`` object as in\n",
        "[here](2_dglgraph.ipynb)). It however contains the information\n",
        "necessary for recovering the original graphs, such as the number of\n",
        "nodes and edges of each graph element.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "50b6d30f",
      "metadata": {
        "id": "50b6d30f"
      },
      "outputs": [],
      "source": [
        "batched_graph, labels = batch\n",
        "print('Number of nodes for each graph element in the batch:', batched_graph.batch_num_nodes())\n",
        "print('Number of edges for each graph element in the batch:', batched_graph.batch_num_edges())\n",
        "\n",
        "# Recover the original graph elements from the minibatch\n",
        "graphs = dgl.unbatch(batched_graph)\n",
        "print('The original graphs in the minibatch:')\n",
        "print(graphs)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "20dd37e3",
      "metadata": {
        "id": "20dd37e3"
      },
      "source": [
        "Define Model\n",
        "------------\n",
        "\n",
        "This tutorial will build a two-layer [Graph Convolutional Network\n",
        "(GCN)](http://tkipf.github.io/graph-convolutional-networks/). Each of\n",
        "its layer computes new node representations by aggregating neighbor\n",
        "information. If you have gone through the\n",
        "[introduction](1_introduction.ipynb), you will notice two\n",
        "differences:\n",
        "\n",
        "-  Since the task is to predict a single category for the *entire graph*\n",
        "   instead of for every node, you will need to aggregate the\n",
        "   representations of all the nodes and potentially the edges to form a\n",
        "   graph-level representation. Such process is more commonly referred as\n",
        "   a *readout*. A simple choice is to average the node features of a\n",
        "   graph with ``dgl.mean_nodes()``.\n",
        "\n",
        "-  The input graph to the model will be a batched graph yielded by the\n",
        "   ``GraphDataLoader``. The readout functions provided by DGL can handle\n",
        "   batched graphs so that they will return one representation for each\n",
        "   minibatch element.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "273c7d54",
      "metadata": {
        "id": "273c7d54"
      },
      "outputs": [],
      "source": [
        "from dgl.nn import GraphConv\n",
        "\n",
        "class GCN(nn.Module):\n",
        "    def __init__(self, in_feats, h_feats, num_classes):\n",
        "        super(GCN, self).__init__()\n",
        "        self.conv1 = GraphConv(in_feats, h_feats)\n",
        "        self.conv2 = GraphConv(h_feats, num_classes)\n",
        "\n",
        "    def forward(self, g, in_feat):\n",
        "        h = self.conv1(g, in_feat)\n",
        "        h = F.relu(h)\n",
        "        h = self.conv2(g, h)\n",
        "        g.ndata['h'] = h\n",
        "        return dgl.mean_nodes(g, 'h')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "da37dd2b",
      "metadata": {
        "id": "da37dd2b"
      },
      "source": [
        "Training Loop\n",
        "-------------\n",
        "\n",
        "The training loop iterates over the training set with the\n",
        "``GraphDataLoader`` object and computes the gradients, just like\n",
        "image classification or language modeling.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c909a4bc",
      "metadata": {
        "id": "c909a4bc"
      },
      "outputs": [],
      "source": [
        "# Create the model with given dimensions\n",
        "model = GCN(dataset.dim_nfeats, 16, dataset.gclasses)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
        "\n",
        "for epoch in range(20):\n",
        "    for batched_graph, labels in train_dataloader:\n",
        "        pred = model(batched_graph, batched_graph.ndata['attr'].float())\n",
        "        loss = F.cross_entropy(pred, labels)\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "num_correct = 0\n",
        "num_tests = 0\n",
        "for batched_graph, labels in test_dataloader:\n",
        "    pred = model(batched_graph, batched_graph.ndata['attr'].float())\n",
        "    num_correct += (pred.argmax(1) == labels).sum().item()\n",
        "    num_tests += len(labels)\n",
        "\n",
        "print('Test accuracy:', num_correct / num_tests)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "43a020e7",
      "metadata": {
        "id": "43a020e7"
      },
      "source": [
        "``DGLDataset`` Object Overview\n",
        "------------------------------\n",
        "\n",
        "Your custom graph dataset should inherit the ``dgl.data.DGLDataset``\n",
        "class and implement the following methods:\n",
        "\n",
        "-  ``__getitem__(self, i)``: retrieve the ``i``-th example of the\n",
        "   dataset. An example often contains a single DGL graph, and\n",
        "   occasionally its label.\n",
        "-  ``__len__(self)``: the number of examples in the dataset.\n",
        "-  ``process(self)``: load and process raw data from disk.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7259fa0d",
      "metadata": {
        "id": "7259fa0d"
      },
      "source": [
        "Creating a Dataset for Node Classification or Link Prediction from CSV\n",
        "----------------------------------------------------------------------\n",
        "\n",
        "A node classification dataset often consists of a single graph, as well\n",
        "as its node and edge features.\n",
        "\n",
        "This tutorial takes a small dataset based on [Zachary’s Karate Club\n",
        "network](https://en.wikipedia.org/wiki/Zachary%27s_karate_club). It\n",
        "contains\n",
        "\n",
        "* A ``members.csv`` file containing the attributes of all\n",
        "  members, as well as their attributes.\n",
        "\n",
        "* An ``interactions.csv`` file\n",
        "  containing the pair-wise interactions between two club members.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "309d0b17",
      "metadata": {
        "id": "309d0b17"
      },
      "outputs": [],
      "source": [
        "import urllib.request\n",
        "import pandas as pd\n",
        "urllib.request.urlretrieve(\n",
        "    'https://data.dgl.ai/tutorial/dataset/members.csv', './members.csv')\n",
        "urllib.request.urlretrieve(\n",
        "    'https://data.dgl.ai/tutorial/dataset/interactions.csv', './interactions.csv')\n",
        "\n",
        "members = pd.read_csv('./members.csv')\n",
        "members.head()\n",
        "\n",
        "interactions = pd.read_csv('./interactions.csv')\n",
        "interactions.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7e5e9255",
      "metadata": {
        "id": "7e5e9255"
      },
      "source": [
        "This tutorial treats the members as nodes and interactions as edges. It\n",
        "takes age as a numeric feature of the nodes, affiliated club as the label\n",
        "of the nodes, and edge weight as a numeric feature of the edges.\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: The original Zachary’s Karate Club network does not have\n",
        "   member ages. The ages in this tutorial are generated synthetically\n",
        "   for demonstrating how to add node features into the graph for dataset\n",
        "   creation.\n",
        "\n",
        "</div>\n",
        "\n",
        "<div class=\"alert alert-info\">\n",
        "    \n",
        "**Note**: In practice, taking age directly as a numeric feature may\n",
        "   not work well in machine learning; strategies like binning or\n",
        "   normalizing the feature would work better. This tutorial directly\n",
        "   takes the values as-is for simplicity.\n",
        "\n",
        "</div>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fa09573b",
      "metadata": {
        "id": "fa09573b"
      },
      "source": [
        "## Understanding DGL Datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ac75c55d",
      "metadata": {
        "id": "ac75c55d"
      },
      "outputs": [],
      "source": [
        "import dgl\n",
        "from dgl.data import DGLDataset\n",
        "import torch\n",
        "import os\n",
        "\n",
        "class KarateClubDataset(DGLDataset):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='karate_club')\n",
        "\n",
        "    def process(self):\n",
        "        nodes_data = pd.read_csv('./members.csv')\n",
        "        edges_data = pd.read_csv('./interactions.csv')\n",
        "        node_features = torch.from_numpy(nodes_data['Age'].to_numpy())\n",
        "        node_labels = torch.from_numpy(nodes_data['Club'].astype('category').cat.codes.to_numpy())\n",
        "        edge_features = torch.from_numpy(edges_data['Weight'].to_numpy())\n",
        "        edges_src = torch.from_numpy(edges_data['Src'].to_numpy())\n",
        "        edges_dst = torch.from_numpy(edges_data['Dst'].to_numpy())\n",
        "\n",
        "        self.graph = dgl.graph((edges_src, edges_dst), num_nodes=nodes_data.shape[0])\n",
        "        self.graph.ndata['feat'] = node_features\n",
        "        self.graph.ndata['label'] = node_labels\n",
        "        self.graph.edata['weight'] = edge_features\n",
        "\n",
        "        # If your dataset is a node classification dataset, you will need to assign\n",
        "        # masks indicating whether a node belongs to training, validation, and test set.\n",
        "        n_nodes = nodes_data.shape[0]\n",
        "        n_train = int(n_nodes * 0.6)\n",
        "        n_val = int(n_nodes * 0.2)\n",
        "        train_mask = torch.zeros(n_nodes, dtype=torch.bool)\n",
        "        val_mask = torch.zeros(n_nodes, dtype=torch.bool)\n",
        "        test_mask = torch.zeros(n_nodes, dtype=torch.bool)\n",
        "        train_mask[:n_train] = True\n",
        "        val_mask[n_train:n_train + n_val] = True\n",
        "        test_mask[n_train + n_val:] = True\n",
        "        self.graph.ndata['train_mask'] = train_mask\n",
        "        self.graph.ndata['val_mask'] = val_mask\n",
        "        self.graph.ndata['test_mask'] = test_mask\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        return self.graph\n",
        "\n",
        "    def __len__(self):\n",
        "        return 1\n",
        "\n",
        "dataset = KarateClubDataset()\n",
        "graph = dataset[0]\n",
        "\n",
        "print(graph)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "852ed071",
      "metadata": {
        "id": "852ed071"
      },
      "source": [
        "Since a link prediction dataset only involves a single graph, preparing\n",
        "a link prediction dataset will have the same experience as preparing a\n",
        "node classification dataset.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "89bdd3da",
      "metadata": {
        "id": "89bdd3da"
      },
      "source": [
        "Creating a Dataset for Graph Classification from CSV\n",
        "----------------------------------------------------\n",
        "\n",
        "Creating a graph classification dataset involves implementing\n",
        "``__getitem__`` to return both the graph and its graph-level label.\n",
        "\n",
        "This tutorial demonstrates how to create a graph classification dataset\n",
        "with the following synthetic CSV data:\n",
        "\n",
        "-  ``graph_edges.csv``: containing three columns:\n",
        "\n",
        "   -  ``graph_id``: the ID of the graph.\n",
        "   -  ``src``: the source node of an edge of the given graph.\n",
        "   -  ``dst``: the destination node of an edge of the given graph.\n",
        "\n",
        "-  ``graph_properties.csv``: containing three columns:\n",
        "\n",
        "   -  ``graph_id``: the ID of the graph.\n",
        "   -  ``label``: the label of the graph.\n",
        "   -  ``num_nodes``: the number of nodes in the graph.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "84f918d7",
      "metadata": {
        "id": "84f918d7"
      },
      "outputs": [],
      "source": [
        "urllib.request.urlretrieve(\n",
        "    'https://data.dgl.ai/tutorial/dataset/graph_edges.csv', './graph_edges.csv')\n",
        "urllib.request.urlretrieve(\n",
        "    'https://data.dgl.ai/tutorial/dataset/graph_properties.csv', './graph_properties.csv')\n",
        "edges = pd.read_csv('./graph_edges.csv')\n",
        "properties = pd.read_csv('./graph_properties.csv')\n",
        "\n",
        "edges.head()\n",
        "\n",
        "properties.head()\n",
        "\n",
        "class SyntheticDataset(DGLDataset):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='synthetic')\n",
        "\n",
        "    def process(self):\n",
        "        edges = pd.read_csv('./graph_edges.csv')\n",
        "        properties = pd.read_csv('./graph_properties.csv')\n",
        "        self.graphs = []\n",
        "        self.labels = []\n",
        "\n",
        "        # Create a graph for each graph ID from the edges table.\n",
        "        # First process the properties table into two dictionaries with graph IDs as keys.\n",
        "        # The label and number of nodes are values.\n",
        "        label_dict = {}\n",
        "        num_nodes_dict = {}\n",
        "        for _, row in properties.iterrows():\n",
        "            label_dict[row['graph_id']] = row['label']\n",
        "            num_nodes_dict[row['graph_id']] = row['num_nodes']\n",
        "\n",
        "        # For the edges, first group the table by graph IDs.\n",
        "        edges_group = edges.groupby('graph_id')\n",
        "\n",
        "        # For each graph ID...\n",
        "        for graph_id in edges_group.groups:\n",
        "            # Find the edges as well as the number of nodes and its label.\n",
        "            edges_of_id = edges_group.get_group(graph_id)\n",
        "            src = edges_of_id['src'].to_numpy()\n",
        "            dst = edges_of_id['dst'].to_numpy()\n",
        "            num_nodes = num_nodes_dict[graph_id]\n",
        "            label = label_dict[graph_id]\n",
        "\n",
        "            # Create a graph and add it to the list of graphs and labels.\n",
        "            g = dgl.graph((src, dst), num_nodes=num_nodes)\n",
        "            self.graphs.append(g)\n",
        "            self.labels.append(label)\n",
        "\n",
        "        # Convert the label list to tensor for saving.\n",
        "        self.labels = torch.LongTensor(self.labels)\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        return self.graphs[i], self.labels[i]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.graphs)\n",
        "\n",
        "dataset = SyntheticDataset()\n",
        "graph, label = dataset[0]\n",
        "print(graph, label)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "CREDITS: Vaibhav Arora, Department of CSE, IIT Bhilai"
      ],
      "metadata": {
        "id": "OZ-shX0eXoHE"
      },
      "id": "OZ-shX0eXoHE"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.11.6"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}