{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Building a Graph Convolutional Network Using Sparse Matrices\n",
        "\n",
        "This tutorial illustrates step-by-step how to write and train a Graph Convolutional Network ([Kipf et al. (2017)](https://arxiv.org/abs/1609.02907)) using DGL's sparse matrix APIs.\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/sparse/gcn.ipynb) [![GitHub](https://img.shields.io/badge/-View%20on%20GitHub-181717?logo=github&logoColor=ffffff)](https://github.com/dmlc/dgl/blob/master/notebooks/sparse/gcn.ipynb)"
      ],
      "metadata": {
        "id": "_iqWrPwxtZr6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Install required packages.\n",
        "import os\n",
        "import torch\n",
        "os.environ['TORCH'] = torch.__version__\n",
        "os.environ['DGLBACKEND'] = \"pytorch\"\n",
        "\n",
        "# Uncomment below to install required packages. If the CUDA version is not 11.6,\n",
        "# check the https://www.dgl.ai/pages/start.html to find the supported CUDA\n",
        "# version and corresponding command to install DGL.\n",
        "#!pip install dgl -f https://data.dgl.ai/wheels/cu118/repo.html > /dev/null\n",
        "\n",
        "try:\n",
        "    import dgl\n",
        "    installed = True\n",
        "except ImportError:\n",
        "    installed = False\n",
        "print(\"DGL installed!\" if installed else \"DGL not found!\")"
      ],
      "metadata": {
        "id": "FTqB360eRvya"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Graph Convolutional Layer\n",
        "\n",
        "Mathematically, the graph convolutional layer is defined as:\n",
        "\n",
        "$$f(X^{(l)}, A) = \\sigma(\\bar{D}^{-\\frac{1}{2}}\\bar{A}\\bar{D}^{-\\frac{1}{2}}X^{(l)}W^{(l)})$$\n",
        "\n",
        "with $\\bar{A} = A + I$, where $A$ denotes the adjacency matrix and $I$ denotes the identity matrix, $\\bar{D}$ refers to the diagonal node degree matrix of $\\bar{A}$ and $W^{(l)}$ denotes a trainable weight matrix. $\\sigma$ refers to a non-linear activation (e.g. relu).\n",
        "\n",
        "The code below shows how to implement it using the `dgl.sparse` package. The core operations are:\n",
        "\n",
        "* `dgl.sparse.identity` creates the identity matrix $I$.\n",
        "* The augmented adjacency matrix $\\bar{A}$ is then computed by adding the identity matrix to the adjacency matrix $A$.\n",
        "* `A_hat.sum(0)` aggregates the augmented adjacency matrix $\\bar{A}$ along the first dimension which gives the degree vector of the augmented graph. The diagonal degree matrix $\\bar{D}$ is then created by `dgl.sparse.diag`.\n",
        "* Compute $\\bar{D}^{-\\frac{1}{2}}$.\n",
        "* `D_hat_invsqrt @ A_hat @ D_hat_invsqrt` computes the convolution matrix which is then multiplied by the linearly transformed node features."
      ],
      "metadata": {
        "id": "r3qB1atg_ld0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "import dgl.sparse as dglsp\n",
        "\n",
        "class GCNLayer(nn.Module):\n",
        "    def __init__(self, in_size, out_size):\n",
        "        super(GCNLayer, self).__init__()\n",
        "        self.W = nn.Linear(in_size, out_size)\n",
        "\n",
        "    def forward(self, A, X):\n",
        "        ########################################################################\n",
        "        # (HIGHLIGHT) Compute the symmetrically normalized adjacency matrix with\n",
        "        # Sparse Matrix API\n",
        "        ########################################################################\n",
        "        I = dglsp.identity(A.shape)\n",
        "        A_hat = A + I\n",
        "        D_hat = dglsp.diag(A_hat.sum(0))\n",
        "        D_hat_invsqrt = D_hat ** -0.5\n",
        "        return D_hat_invsqrt @ A_hat @ D_hat_invsqrt @ self.W(X)"
      ],
      "metadata": {
        "id": "Y4I4EhHQ_kKb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "A Graph Convolutional Network is then defined by stacking this layer."
      ],
      "metadata": {
        "id": "bvP7O2IwV_c7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a GCN with the GCN layer.\n",
        "class GCN(nn.Module):\n",
        "    def __init__(self, in_size, out_size, hidden_size):\n",
        "        super(GCN, self).__init__()\n",
        "        self.conv1 = GCNLayer(in_size, hidden_size)\n",
        "        self.conv2 = GCNLayer(hidden_size, out_size)\n",
        "\n",
        "    def forward(self, A, X):\n",
        "        X = self.conv1(A, X)\n",
        "        X = F.relu(X)\n",
        "        return self.conv2(A, X)"
      ],
      "metadata": {
        "id": "BHX3vRjDWJTO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training the GCN\n",
        "\n",
        "We then train the GCN model on the Cora dataset for node classification. Note that since the model expects an adjacency matrix as the first argument, we first construct the adjacency matrix from the graph using the `dgl.sparse.from_coo` API which returns a DGL `SparseMatrix` object."
      ],
      "metadata": {
        "id": "2Qw7fTdGNnEp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(g, pred):\n",
        "    label = g.ndata[\"label\"]\n",
        "    val_mask = g.ndata[\"val_mask\"]\n",
        "    test_mask = g.ndata[\"test_mask\"]\n",
        "\n",
        "    # Compute accuracy on validation/test set.\n",
        "    val_acc = (pred[val_mask] == label[val_mask]).float().mean()\n",
        "    test_acc = (pred[test_mask] == label[test_mask]).float().mean()\n",
        "    return val_acc, test_acc\n",
        "\n",
        "def train(model, g):\n",
        "    features = g.ndata[\"feat\"]\n",
        "    label = g.ndata[\"label\"]\n",
        "    train_mask = g.ndata[\"train_mask\"]\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)\n",
        "    loss_fcn = nn.CrossEntropyLoss()\n",
        "\n",
        "    # Preprocess to get the adjacency matrix of the graph.\n",
        "    indices = torch.stack(g.edges())\n",
        "    N = g.num_nodes()\n",
        "    A = dglsp.spmatrix(indices, shape=(N, N))\n",
        "\n",
        "    for epoch in range(100):\n",
        "        model.train()\n",
        "\n",
        "        # Forward.\n",
        "        logits = model(A, features)\n",
        "\n",
        "        # Compute loss with nodes in the training set.\n",
        "        loss = loss_fcn(logits[train_mask], label[train_mask])\n",
        "\n",
        "        # Backward.\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        # Compute prediction.\n",
        "        pred = logits.argmax(dim=1)\n",
        "\n",
        "        # Evaluate the prediction.\n",
        "        val_acc, test_acc = evaluate(g, pred)\n",
        "        if epoch % 5 == 0:\n",
        "            print(\n",
        "                f\"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}\"\n",
        "                f\", test acc: {test_acc:.3f}\"\n",
        "            )\n",
        "\n",
        "\n",
        "# Load graph from the existing dataset.\n",
        "dataset = dgl.data.CoraGraphDataset()\n",
        "g = dataset[0]\n",
        "\n",
        "# Create model.\n",
        "feature = g.ndata['feat']\n",
        "in_size = feature.shape[1]\n",
        "out_size = dataset.num_classes\n",
        "gcn_model = GCN(in_size, out_size, 16)\n",
        "\n",
        "# Kick off training.\n",
        "train(gcn_model, g)"
      ],
      "metadata": {
        "id": "5Sp1B1_QHgC2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "552e2c22-44f4-4495-c7f9-a57f13484270"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading /root/.dgl/cora_v2.zip from https://data.dgl.ai/dataset/cora_v2.zip...\n",
            "Extracting file to /root/.dgl/cora_v2\n",
            "Finished data loading and preprocessing.\n",
            "  NumNodes: 2708\n",
            "  NumEdges: 10556\n",
            "  NumFeats: 1433\n",
            "  NumClasses: 7\n",
            "  NumTrainingSamples: 140\n",
            "  NumValidationSamples: 500\n",
            "  NumTestSamples: 1000\n",
            "Done saving data into cached files.\n",
            "In epoch 0, loss: 1.954, val acc: 0.114, test acc: 0.103\n",
            "In epoch 5, loss: 1.921, val acc: 0.158, test acc: 0.147\n",
            "In epoch 10, loss: 1.878, val acc: 0.288, test acc: 0.283\n",
            "In epoch 15, loss: 1.822, val acc: 0.344, test acc: 0.353\n",
            "In epoch 20, loss: 1.751, val acc: 0.388, test acc: 0.389\n",
            "In epoch 25, loss: 1.663, val acc: 0.406, test acc: 0.410\n",
            "In epoch 30, loss: 1.562, val acc: 0.472, test acc: 0.481\n",
            "In epoch 35, loss: 1.450, val acc: 0.558, test acc: 0.573\n",
            "In epoch 40, loss: 1.333, val acc: 0.636, test acc: 0.641\n",
            "In epoch 45, loss: 1.216, val acc: 0.684, test acc: 0.683\n",
            "In epoch 50, loss: 1.102, val acc: 0.726, test acc: 0.713\n",
            "In epoch 55, loss: 0.996, val acc: 0.740, test acc: 0.740\n",
            "In epoch 60, loss: 0.899, val acc: 0.754, test acc: 0.760\n",
            "In epoch 65, loss: 0.813, val acc: 0.762, test acc: 0.771\n",
            "In epoch 70, loss: 0.737, val acc: 0.768, test acc: 0.781\n",
            "In epoch 75, loss: 0.671, val acc: 0.776, test acc: 0.786\n",
            "In epoch 80, loss: 0.614, val acc: 0.784, test acc: 0.790\n",
            "In epoch 85, loss: 0.566, val acc: 0.780, test acc: 0.788\n",
            "In epoch 90, loss: 0.524, val acc: 0.780, test acc: 0.791\n",
            "In epoch 95, loss: 0.489, val acc: 0.772, test acc: 0.795\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Check out the full example script* [here](https://github.com/dmlc/dgl/blob/master/examples/sparse/gcn.py)."
      ],
      "metadata": {
        "id": "yQnJZvE9ZduM"
      }
    }
  ]
}