SparceArray_Graph_Networks_ipynb.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/zwang86/DLSys-Fall-2022-Report/blob/main/SparceArray_Graph_Networks_ipynb.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Introduction\n",
        "\n",
        "We extend the needle library by adding sparse matrices and related operations and autodiff. It supports both CPU and CUDA backends. On top of that, we build the Graph Convolutional Network and the Graph Attention Network. We train and compare these two networks and a baseline fully-connected network on the Cora dataset, and show that graph networks can achieve higher accuracy without too much performance cost."
      ],
      "metadata": {
        "id": "NZlmFsAanJUc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Environment Setup\n",
        "\n"
      ],
      "metadata": {
        "id": "rtr_6FXOWLna"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The project is available in [Github repository](https://github.com/MashPlant/dlsys-proj). You may need to request access first.\n",
        "You can also access the code directly from [shared directory](https://drive.google.com/drive/folders/1bu0THsbS5IuQzBHvFPdmy6qdbcuNTAS7?usp=share_link). "
      ],
      "metadata": {
        "id": "ymHS4cUF28bJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Mount Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "%cd /content/drive/MyDrive/\n",
        "!mkdir -p 10714\n",
        "%cd /content/drive/MyDrive/10714"
      ],
      "metadata": {
        "id": "Ti9NwkTT3Mvt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Configure Github and repository\n",
        "!git config --global user.email \"{YOUR_EMAIL}\"\n",
        "!git config --global user.name \"{YOUR_USERNAME}\"\n",
        "!git clone https://\"{GITHUB_TOKEN}\"@github.com/MashPlant/dlsys-proj.git\n",
        "%cd /content/drive/MyDrive/10714/dlsys-proj\n",
        "!git pull"
      ],
      "metadata": {
        "id": "dMH7kEbAPW7q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Add Python PATH\n",
        "import sys\n",
        "sys.path.append('./python')"
      ],
      "metadata": {
        "id": "zpTAPqqKPeDA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Install pybind11\n",
        "!pip3 install pybind11"
      ],
      "metadata": {
        "id": "U6XDeOn6GfpD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Build CPU and CUDA backend\n",
        "!make"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0WJm06QkPfRf",
        "outputId": "61f66959-f555-4fe9-d1ab-d4ba379a8063"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "-- Found pybind11: /usr/local/lib/python3.8/dist-packages/pybind11/include (found version \"2.10.1\")\n",
            "-- Found cuda, building cuda backend\n",
            "Sun Dec  4 20:48:56 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   58C    P0    29W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n",
            "-- Autodetected CUDA architecture(s):  7.5\n",
            "-- Configuring done\n",
            "-- Generating done\n",
            "-- Build files have been written to: /content/drive/MyDrive/10714/dlsys-proj/build\n",
            "make[1]: Entering directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "make[2]: Entering directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "make[3]: Entering directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "\u001b[35m\u001b[1mConsolidate compiler generated dependencies of target ndarray_backend_cpu\u001b[0m\n",
            "make[3]: Leaving directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "[  0%] Built target ndarray_backend_cpu\n",
            "make[3]: Entering directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "[ 25%] \u001b[34m\u001b[1mBuilding NVCC (Device) object CMakeFiles/ndarray_backend_cuda.dir/src/ndarray_backend_cuda_generated_ndarray_backend_cuda.cu.o\u001b[0m\n",
            "make[3]: Leaving directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "make[3]: Entering directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "[ 50%] \u001b[32m\u001b[1mLinking CXX shared module ../python/needle/backend_ndarray/ndarray_backend_cuda.cpython-38-x86_64-linux-gnu.so\u001b[0m\n",
            "make[3]: Leaving directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "[ 50%] Built target ndarray_backend_cuda\n",
            "make[2]: Leaving directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n",
            "make[1]: Leaving directory '/content/drive/MyDrive/10714/dlsys-proj/build'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Sparse Matrix Operations\n",
        "\n",
        "## Construction and Conversion\n",
        "\n",
        "For simplicity, we only support 2-D sparse matrices instead of general multidimensional sparse tensors, considering their usage in the project. There are many data formats to represent sparse matrices, such as DOK, COO, LIL, CSR, etc. We choose to use the CSR format. The Compressed Sparse Row (CSR) format is a common way of storing sparse matrices in memory. A matrix is represented as three arrays:\n",
        "\n",
        "- `row`: Length is the number of rows plus one. It marks the starting and ending positions of a row in the `col` and `val` array. Specifically, the elements in row `r` are stored in `col[row[r]..row[r + 1]]` and `val[row[r]..row[r + 1]]` (left-closed and right-open interval).\n",
        "- `col`: Length is the number of non-zero elements. It contains the column indices of the non-zero elements.\n",
        "- `val`: Length is the number of non-zero elements. It contains the values of the non-zero elements.\n",
        "\n",
        "It has some advantages that we value, including less memory usage (only 8 bytes for each `float32` non-zero element) and a hierarchical structure that can enable parallelism (especially for CUDA in the project).\n",
        "\n",
        "We do not modify the `NDArray` class from the original `needle` library. Instead, we create a `SparseArray` class with a similar status with `NDArray`, which means that `Value.cached_data` can be either a `NDArray` or a `SparseArray`.\n",
        "We export a set of functions with `sparse` as the name prefix in CPU and CUDA backends. They are used in the methods in `SparseArray`, and are mostly straightforward.\n",
        "\n",
        "We provide two methods to convert a sparse matrix from/to a dense array, which is represented as a contiguous piece of memory.\n",
        "\n",
        "- `sparse_from_dense`: Convert a dense array to a `SparceArray` instance. It iterates over all the elements and only keeps the non-zero ones.\n",
        "- `sparse_to_dense`: Convert a `SparseArray` instance to a dense array. It zeros the ouput array and fill the non-zero elements.\n",
        "\n",
        "## Operations\n",
        "\n",
        "We also define a similar set of methods in `SparseArray` as `NDArray`. Since the operations in `ops.py` mostly just forward the call to the underlying array class, we can reuse most of them, and only need to do very few modifications.\n",
        "\n",
        "There is a slight difference in implementation. The operations of `NDArray` normally first allocates an `NDArray` instance, and then stores the calculation result in it. On the other hand, the number of non-zero elements of the resulting `SparseArray` may depend on the specific operation, and we cannot allocate it in advance, so we let these functions return `SparseArray` directly.\n",
        "\n",
        "In order to support both CPU and CUDA backends, for most operations, we can simply replace the outermost one or two loops of the CPU version with CUDA blocks or block+threads.\n",
        "For some operations (such as `sparse_gat_grad`), doing this directly will cause data racing, and we need to modify the computation order to adapt to GPU.\n",
        "For some operations (such as `sparse_transpose`), it is too hard to implement them with GPU, and we choose to transfer the data to CPU, invoke the CPU version, and transfer the data to GPU.\n",
        "\n",
        "- `sparse_[add|sub|mul|div|rdiv]_[scalar|dense|sparse]` (in total 15 functions): Perform basic scalar or elemen-wise operations on sparse-scalar, sparse-dense and sparse-sparse operands. For sparse-sparse, we only support the case where the non-zero element locations of the two operands are the same. This is enough for our usage.\n",
        "- `sparse_dense_matmul`: Perform matrix multiplication on a sparse matrix and a dense matrix. Suppose the two operands have shape `m x n` and `n x k` respectively, then the dense-dense matmul takes `O(mnk)` time, while the sparse-dense matmul takes `O((number of non-zero elements in A) x k)` time. Since the edge matrix is sparse in our usage, it is very likely to reduce the cubic time into quadratic time.\n",
        "- `sparse_transpose`: The transposition of CSR matrices is a little tricky to implement efficiently. By definition, the `row` array in CSR maintains an exclusive prefix sum of the number of elements in each row. We first count the number of elements in each column (because of transposition), and then compute an inclusive prefix sum of the counters. We then visits elements in the original matrix. By using this inclusive prefix sum and modifying it in the fly, we can place the elements correctly and get the exclusive prefix sum. We think that this is too complex for a GPU implementation.  \n",
        "- `sparse_gat`: We will elaborate this in the GAT section.\n",
        "- `sparse_gat_grad_from_[sparse|dense]`: The gradient of the `sparse_gat`. It takes the upstream gradient and calculates the operands' final gradients directly (i.e, the caller doesn't need to do the chain-rule multiplication). The `[sparse|dense]` variants are explained in the following section.\n",
        "- `sparse_softmax`: The softmax function for sparse matrices on axis-1. \n",
        "- `sparse_softmax_grad_from_[sparse|dense]`: The gradient of `sparse_softmax`, similar to `sparse_gat_grad_from_[sparse|dense]`.\n",
        "\n",
        "## Autograd\n",
        "\n",
        "It is mostly naturally supported because of the reuse in `ops.py`. However, there is an important difference in `sparse_gat` and `sparse_softmax`. Ideally, we should implement a complex operation by decomposing it into several differentiable simple operations. However, we feel it too complex to do so for them. Instead, we implement the computation and gradient logic directly in the backend. Doing so disables us from taking the second-order derivative even though the operation is second-order differentiable, but this has no impact on our usage. Anouther resulting problem is that the upstream gradient matrix may be dense or sparse. We could convert them to one format and handle only that format, but we think it has a performance overhead. So we implement the `[sparse|dense]` variants (with some C++ template tricks, this hardly increases the amount of code).\n",
        "\n",
        "## Modules\n",
        "\n",
        "For the modules, we add a few new modules to support the DNN we want to build in this project, including `GraphConvolution`, `GraphAttention`, and `MultiHeadGAT`. Plus 20 modules we implemented in the previous homework, the total of 23 modules have now supported both dense and sparse matrices and be able to make some powerful DNNs.\n"
      ],
      "metadata": {
        "id": "5qBfNI-tlrTG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tests\n",
        "\n",
        "Note that each test includes many shapes. We did not make `shape` a parameter of testing like our homework, as it would produce too many tests."
      ],
      "metadata": {
        "id": "2CkRuZcrbiaG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# A few test on our new ops with cpu/cuda backend \n",
        "!python3 -m pytest -l -v -k \"test_sparse\""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NoUSa_LkXIYT",
        "outputId": "d2c6eee3-6d4f-43dd-dd5e-5682e3140546"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m============================= test session starts ==============================\u001b[0m\n",
            "platform linux -- Python 3.8.15, pytest-3.6.4, py-1.11.0, pluggy-0.7.1 -- /usr/bin/python3\n",
            "cachedir: .pytest_cache\n",
            "rootdir: /content/drive/MyDrive/10714/dlsys-proj, inifile:\n",
            "plugins: typeguard-2.7.1\n",
            "collected 24 items                                                             \u001b[0m\n",
            "\n",
            "tests/test_sparse.py::test_sparse_init[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                   [  4%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_init[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                   [  8%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_matmul_forward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m         [ 12%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_matmul_forward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m         [ 16%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_matmul_backward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m        [ 20%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_matmul_backward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m        [ 25%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_add[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 29%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_add[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 33%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_sub[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 37%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_sub[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 41%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_mul[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 45%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_mul[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 50%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_div[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 54%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_div[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                    [ 58%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_transpose_forward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m      [ 62%]\u001b[0m\n",
            "tests/test_sparse.py::test_sparse_transpose_forward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m      [ 66%]\u001b[0m\n",
            "tests/test_sparse.py::test_gat_forward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                   [ 70%]\u001b[0m\n",
            "tests/test_sparse.py::test_gat_forward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                   [ 75%]\u001b[0m\n",
            "tests/test_sparse.py::test_gat_backward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m                  [ 79%]\u001b[0m\n",
            "tests/test_sparse.py::test_gat_backward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m                  [ 83%]\u001b[0m\n",
            "tests/test_sparse.py::test_softmax_forward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m               [ 87%]\u001b[0m\n",
            "tests/test_sparse.py::test_softmax_forward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m               [ 91%]\u001b[0m\n",
            "tests/test_sparse.py::test_softmax_backward[device0] \u001b[32mPASSED\u001b[0m\u001b[36m              [ 95%]\u001b[0m\n",
            "tests/test_sparse.py::test_softmax_backward[device1] \u001b[32mPASSED\u001b[0m\u001b[36m              [100%]\u001b[0m\n",
            "\n",
            "\u001b[32m\u001b[1m========================== 24 passed in 53.75 seconds ==========================\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Graph Convolutional Networks\n",
        "\n",
        "Our implementation is based on the paper \"[Simplifying Graph Convolutional Networks](https://arxiv.org/abs/1902.07153)\" published in 2019 (although we didn't do the \"simplifying\" part, we find its illustration for the original GCNs easy to follow).\n",
        "\n",
        "Graph Convolutional Networks (GCNs) are a type of neural network that can process data represented as a graph. They are a generalization of convolutional networks, which are commonly used for image processing. Like convolutional networks, GCNs use convolutional layers to extract local features from the graph structure. The locality in the figure is reflected in the neighborhood relationship within certain hops. As we will describe, it is actually done may matrix multiplication instead of the convolution operator.\n",
        "\n",
        "The GCN operates on a graph $G$ with adjacency matrix $E$ and a set of $n$ nodes $V$. $A$ is a sparse matrix where $e_{ij}$ denotes the edge weight between nodes $v_i$ and $v_j$ (or simply 0/1 to denote whether an edge exists). The spirit of a graph convolution layer is to pass the information of a node's neighbors to the node. Suppose the input of a layer is $H^{(k-1)} \\in \\mathbf{R}^{n,f}$, where there is a feature vector of length $f$ associated with each node. We compute $H^{(k)} = \\sigma(E H^{(k-1)} W)$ as the output of the layer. First, $H^{(k-1)} W$ can be considered as a linear preprocessing layer (can also be extended with a bias term). And then, we multiply it with $A$, which means that the new feature vector of a node will be a (weighted) sum of all the transformed feature vectors of its neighbors. At last, we apply an activation function $\\sigma$.\n",
        "\n",
        "```\n",
        "\n",
        "in_features -> self.preprocess -> hidden_features * num_heads -> self.gn_layers -> hidden_features * num_heads -> self.postprocess -> out_features\n",
        "\n",
        "where:\n",
        "  1. self.preprocess is a nn.Linear layer with in_features and hidden_features * num_heads as input and output, respectively\n",
        "  2. self.gn_layers is a list of nn.MultiHeadGAT layers with hidden_features * num_heads as input and hidden_features * num_heads as output\n",
        "  3. self.postprocess is a nn.Linear layer with hidden_features * num_heads as input and out_features as output\n",
        "```"
      ],
      "metadata": {
        "id": "TwwrGAISmBW-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Graph Attention Networks\n",
        "\n",
        "Our implementation is based on the paper \"[Graph Attention Networks](https://arxiv.org/abs/1710.10903)\" published in 2017. The paper introduces attention mechanisms to the graph structure. Graph Attention Networks (GATs) enhance the original GCNs by allowing the network to learn the importance of different nodes and edges and focus on the most relevant parts when making predictions. Recall that in the original GCNs described above, all neighborhoods contribute to a node with fixed weights (in the case where there are no prescribed edge weights, it is actually a fixed and equal weight), which may not be desirable because they are inherently different.\n",
        "\n",
        "We will use the same graph notation (adjacency matrix $E$ and nodes $V$) in GAT as in GCN. The attention coefficients in a GAT layer are computed using the following formula:\n",
        "\n",
        "$$b_{ij} = f (H_i, H_j)$$\n",
        "\n",
        "$b_{ij}$ is only calucalted if node $j$ is a neighbor of node $i$ (i.e, $e_{ij} \\ne 0$), so it is also sparse. It indicates the importance of node $j$'s features to node $i$. $f$ is a $\\mathbf{R}^{f} \\times \\mathbf{R}^{f} \\rightarrow \\mathbf{R}$ function, and $H_i$ and $H_j$ are the node features of nodes $i$ and $j$ in the input features. The paper suggests that $f$ can be a single-layer feedforward neural network, which means that we maintain a learnable weight vector $\\vec{a} \\in \\mathbf{R}^{2f}$, and $f (H_i, H_j) = \\sigma(\\vec{a}^T (H_i \\parallel H_j))$ ($\\parallel$ denotes concatenation). The activation function $\\sigma$ is further picked as $\\text{LeakyReLU}$.\n",
        "\n",
        "The coefficients $b_{ij}$ are then used to compute the attention weights:\n",
        "\n",
        "$$a_{ij} = \\text{softmax}_j(b_{ij}) = \\frac{\\exp(b_{ij})}{\\sum_{k, e_{ik} \\ne 0} \\exp(b_{ik})} = \\frac{\\exp(\\text{LeakyReLU}(\\vec{a}^T (H_i \\parallel H_j)))}{\\sum_{k, e_{ik} \\ne 0} \\exp(\\text{LeakyReLU}(\\vec{a}^T (H_i \\parallel H_k)))}$$\n",
        "\n",
        "<p align=\"center\">\n",
        "  <img src=\"https://raw.github.com/zwang86/DLSys-Fall-2022-Report/main/img/attention_weights.png\" alt=\"Attention Weights\"/>\n",
        "</p>\n",
        "\n",
        "The attention weights are then used to weigh the contributions of different nodes to the output. Let $A$ denote the attention weight matrix. The output of the GAT layer is computed as $H^{(k)} = \\sigma (A H^{(k-1)} W)$ (the only difference between this equation and that of GCN is that the adjacent matrix $E$ is replaced with $A$).\n",
        "\n",
        "In our implementation, `sparse_gat` and `sparse_softmax` are especially used to build the GAT layer. Simply put it, `sparse_gat` computes the $b_{ij}$ described above, and `sparse_softmax` is just a standard softmax over axis-1. `sparse_gat_grad_[sparse|dense]` only computes gradients for $H$ and $\\vec{a}$, but not $E$. It makes sense because $E$ is an input graph property and doesn't require gradients.\n",
        "\n",
        "We also implemented the multi-head attention structure in the paper in the `MultiHeadGAT` module. It takes the output of several GAT layers to stabilize the learning process. We only implemented the concatenation policy, which can be represented as\n",
        "\n",
        "$$H^{(k)} = \\parallel_{k=1}^K \\sigma (A^{(k)} H^{(k-1)} W^{(k)})$$\n"
      ],
      "metadata": {
        "id": "6o2HCCZcmGLV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Models and Training\n",
        "\n",
        "We implemented the Graph Convolution Model and the Graph Attention Model, together with a baseline fully-connected model, in `apps/gn.py`. The baseline model have roughly the same number of parameters as the Graph Convolution Model.\n",
        "\n",
        "<p align=\"center\">\n",
        "  <img src=\"https://raw.github.com/zwang86/DLSys-Fall-2022-Report/main/img/graph_networks.png\" alt=\"Graph Networks\"/>\n",
        "</p>\n",
        "\n",
        "\n",
        "```\n",
        "> python3 apps/gn.py [device] [model]\n",
        "```\n",
        "\n",
        "- device: cuda, cpu\n",
        "- model: baseline, gc, gat\n",
        "\n",
        "We evaluate the models on the Cora dataset. It a commonly used benchmark dataset for evaluating graph-based machine learning algorithms, on tasks such as node classification and link prediction. It consists of 2708 scientific publications classified into one of seven classes. The dataset includes the content (a feature vector) of the papers, the citation relationship among the papers, and the class labels for each paper. The dataset is already included in our repository under `data/`. You can also download it from https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz.\n",
        "\n",
        "It is not trivial to split the train/validation/test sets on the graph data because data points are not independent.\n",
        "We may apply graph networks in the *transductive* setting, which means that the sets are on the same graph.\n",
        "The entire graph and all features (including validation/test features) are visible during training. Only validation/test labels are invisible.\n",
        "There are many ways to handle this setting, and we choose to introduce a new network layer: `MaskedSoftmaxLoss`. It uses a mask so that only the train/validation output v.s. true labels are taken into loss computation.\n",
        "Note that this is not the general case, and these sets may be on independent graphs in other applications of graph networks.\n"
      ],
      "metadata": {
        "id": "upFggTIjU2VG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Baseline Result\n",
        "\n",
        "The output format for all the results is `train = {train accuracy}, {train loss}, val = {validation accuracy}, {validation loss}`."
      ],
      "metadata": {
        "id": "LFA_fZu5QZCB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python3 apps/gn.py cuda baseline"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2Qola37oQcbo",
        "outputId": "8267ca87-b1a1-4002-b30a-5b78455232e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "---Selecting backend:  cuda() ---\n",
            "---Training baseline model---\n",
            "epoch 1, elapsed 0.03s: train = 0.2950743494423792, 3.8048384189605713, val = 0.329136690647482, 3.6662778854370117\n",
            "epoch 2, elapsed 0.04s: train = 0.29460966542750927, 2.0384480953216553, val = 0.3057553956834532, 2.073040008544922\n",
            "epoch 3, elapsed 0.06s: train = 0.41821561338289964, 1.5862475633621216, val = 0.32194244604316546, 1.7222514152526855\n",
            "epoch 4, elapsed 0.08s: train = 0.5455390334572491, 1.3163700103759766, val = 0.42805755395683454, 1.5298867225646973\n",
            "epoch 5, elapsed 0.09s: train = 0.6979553903345725, 1.0869649648666382, val = 0.5413669064748201, 1.3500683307647705\n",
            "epoch 6, elapsed 0.11s: train = 0.7727695167286245, 0.8974674940109253, val = 0.6205035971223022, 1.1971731185913086\n",
            "epoch 7, elapsed 0.12s: train = 0.8336431226765799, 0.7194642424583435, val = 0.697841726618705, 1.0586415529251099\n",
            "epoch 8, elapsed 0.14s: train = 0.8698884758364313, 0.5696262717247009, val = 0.7212230215827338, 0.9484922289848328\n",
            "epoch 9, elapsed 0.15s: train = 0.8977695167286245, 0.4586840569972992, val = 0.7302158273381295, 0.8803451061248779\n",
            "epoch 10, elapsed 0.17s: train = 0.9131040892193308, 0.37842002511024475, val = 0.7212230215827338, 0.8445652723312378\n",
            "epoch 11, elapsed 0.18s: train = 0.9228624535315985, 0.3140508234500885, val = 0.7302158273381295, 0.8243606090545654\n",
            "epoch 12, elapsed 0.19s: train = 0.9395910780669146, 0.25728490948677063, val = 0.7392086330935251, 0.8081569075584412\n",
            "epoch 13, elapsed 0.21s: train = 0.949814126394052, 0.20830726623535156, val = 0.7482014388489209, 0.7952264547348022\n",
            "epoch 14, elapsed 0.23s: train = 0.9614312267657993, 0.16851919889450073, val = 0.7535971223021583, 0.7905815243721008\n",
            "epoch 15, elapsed 0.25s: train = 0.9730483271375465, 0.13756045699119568, val = 0.7589928057553957, 0.7951352000236511\n",
            "epoch 16, elapsed 0.26s: train = 0.9772304832713755, 0.1140422523021698, val = 0.7643884892086331, 0.808445155620575\n",
            "epoch 17, elapsed 0.28s: train = 0.9823420074349443, 0.09569018334150314, val = 0.7643884892086331, 0.8270772695541382\n",
            "epoch 18, elapsed 0.31s: train = 0.9869888475836431, 0.0810079574584961, val = 0.7661870503597122, 0.8490467667579651\n",
            "epoch 19, elapsed 0.32s: train = 0.9888475836431226, 0.06915053725242615, val = 0.762589928057554, 0.8722792863845825\n",
            "epoch 20, elapsed 0.33s: train = 0.9897769516728625, 0.059608928859233856, val = 0.7661870503597122, 0.8958348035812378\n",
            "epoch 21, elapsed 0.35s: train = 0.9902416356877324, 0.05196094140410423, val = 0.7643884892086331, 0.9195823669433594\n",
            "epoch 22, elapsed 0.36s: train = 0.9911710037174721, 0.045706164091825485, val = 0.762589928057554, 0.9436001181602478\n",
            "epoch 23, elapsed 0.37s: train = 0.991635687732342, 0.040542881935834885, val = 0.7607913669064749, 0.9674535393714905\n",
            "epoch 24, elapsed 0.38s: train = 0.9925650557620818, 0.036157138645648956, val = 0.7589928057553957, 0.9908126592636108\n",
            "epoch 25, elapsed 0.40s: train = 0.9930297397769516, 0.03238988295197487, val = 0.7571942446043165, 1.0134981870651245\n",
            "epoch 26, elapsed 0.41s: train = 0.9948884758364313, 0.029132219031453133, val = 0.7553956834532374, 1.0351248979568481\n",
            "epoch 27, elapsed 0.42s: train = 0.9962825278810409, 0.026293635368347168, val = 0.7535971223021583, 1.0556797981262207\n",
            "epoch 28, elapsed 0.43s: train = 0.9967472118959108, 0.023795736953616142, val = 0.7517985611510791, 1.07516610622406\n",
            "epoch 29, elapsed 0.45s: train = 0.9976765799256505, 0.02158302068710327, val = 0.7535971223021583, 1.0935781002044678\n",
            "epoch 30, elapsed 0.46s: train = 0.9976765799256505, 0.01966165006160736, val = 0.75, 1.1107500791549683\n",
            "epoch 31, elapsed 0.47s: train = 0.9976765799256505, 0.017979294061660767, val = 0.7464028776978417, 1.1267403364181519\n",
            "epoch 32, elapsed 0.49s: train = 0.9976765799256505, 0.016506502404808998, val = 0.7517985611510791, 1.1415750980377197\n",
            "epoch 33, elapsed 0.50s: train = 0.9976765799256505, 0.015206918120384216, val = 0.7535971223021583, 1.1555073261260986\n",
            "epoch 34, elapsed 0.51s: train = 0.9976765799256505, 0.014060746878385544, val = 0.7535971223021583, 1.1686573028564453\n",
            "epoch 35, elapsed 0.53s: train = 0.9986059479553904, 0.013046237640082836, val = 0.7517985611510791, 1.1811447143554688\n",
            "epoch 36, elapsed 0.54s: train = 0.9986059479553904, 0.012135262601077557, val = 0.7553956834532374, 1.193103551864624\n",
            "epoch 37, elapsed 0.57s: train = 0.9986059479553904, 0.011315506882965565, val = 0.7535971223021583, 1.2044037580490112\n",
            "epoch 38, elapsed 0.58s: train = 0.9986059479553904, 0.010583366267383099, val = 0.7553956834532374, 1.2149546146392822\n",
            "epoch 39, elapsed 0.59s: train = 0.9986059479553904, 0.00992431491613388, val = 0.7553956834532374, 1.225098967552185\n",
            "epoch 40, elapsed 0.61s: train = 0.9986059479553904, 0.00932861678302288, val = 0.7553956834532374, 1.2347638607025146\n",
            "---Total training time: 0.6086337566375732 seconds ---\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Graph Network Result"
      ],
      "metadata": {
        "id": "IZ0Zs9hiQiwY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python3 apps/gn.py cuda gc"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "229c0664-1d4c-4875-d9f6-1ca3f6a1c093",
        "id": "hloSZ57nZjZh"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "---Selecting backend:  cuda() ---\n",
            "---Training graph convolution network---\n",
            "epoch 1, elapsed 0.02s: train = 0.29275092936802977, 6.315654754638672, val = 0.329136690647482, 5.342335224151611\n",
            "epoch 2, elapsed 0.05s: train = 0.1570631970260223, 6.225630760192871, val = 0.1618705035971223, 6.348045825958252\n",
            "epoch 3, elapsed 0.08s: train = 0.22072490706319703, 3.554532766342163, val = 0.21402877697841727, 3.7679293155670166\n",
            "epoch 4, elapsed 0.11s: train = 0.2578996282527881, 2.47625994682312, val = 0.17805755395683454, 2.69366455078125\n",
            "epoch 5, elapsed 0.13s: train = 0.266728624535316, 2.02108097076416, val = 0.210431654676259, 2.162968397140503\n",
            "epoch 6, elapsed 0.16s: train = 0.5004646840148699, 1.476326823234558, val = 0.4658273381294964, 1.5470199584960938\n",
            "epoch 7, elapsed 0.19s: train = 0.5013940520446096, 1.3853263854980469, val = 0.5251798561151079, 1.353360652923584\n",
            "epoch 8, elapsed 0.21s: train = 0.5204460966542751, 1.3371285200119019, val = 0.5359712230215827, 1.3054759502410889\n",
            "epoch 9, elapsed 0.23s: train = 0.5613382899628253, 1.2182918787002563, val = 0.5683453237410072, 1.2310770750045776\n",
            "epoch 10, elapsed 0.26s: train = 0.616635687732342, 1.105104923248291, val = 0.6025179856115108, 1.1544663906097412\n",
            "epoch 11, elapsed 0.29s: train = 0.6635687732342007, 1.0089702606201172, val = 0.6474820143884892, 1.0821284055709839\n",
            "epoch 12, elapsed 0.31s: train = 0.6872676579925651, 0.9491685032844543, val = 0.6654676258992805, 1.039381504058838\n",
            "epoch 13, elapsed 0.33s: train = 0.7095724907063197, 0.9070488214492798, val = 0.6600719424460432, 1.0088797807693481\n",
            "epoch 14, elapsed 0.36s: train = 0.7174721189591078, 0.8683772683143616, val = 0.6690647482014388, 0.9745235443115234\n",
            "epoch 15, elapsed 0.38s: train = 0.7295539033457249, 0.82168048620224, val = 0.685251798561151, 0.9222384691238403\n",
            "epoch 16, elapsed 0.41s: train = 0.7434944237918215, 0.7777501940727234, val = 0.6960431654676259, 0.8704046607017517\n",
            "epoch 17, elapsed 0.43s: train = 0.7560408921933085, 0.7421087622642517, val = 0.7122302158273381, 0.8307932615280151\n",
            "epoch 18, elapsed 0.45s: train = 0.7685873605947955, 0.710239589214325, val = 0.7320143884892086, 0.7993310689926147\n",
            "epoch 19, elapsed 0.47s: train = 0.7820631970260223, 0.6787621974945068, val = 0.7482014388489209, 0.7727622985839844\n",
            "epoch 20, elapsed 0.49s: train = 0.7932156133828996, 0.6481716632843018, val = 0.7571942446043165, 0.7504745721817017\n",
            "epoch 21, elapsed 0.52s: train = 0.8057620817843866, 0.6183070540428162, val = 0.7661870503597122, 0.7320624589920044\n",
            "epoch 22, elapsed 0.55s: train = 0.8178438661710037, 0.5885531902313232, val = 0.7715827338129496, 0.717424213886261\n",
            "epoch 23, elapsed 0.57s: train = 0.828996282527881, 0.5606794953346252, val = 0.7805755395683454, 0.7051843404769897\n",
            "epoch 24, elapsed 0.59s: train = 0.8382899628252788, 0.5343722105026245, val = 0.7823741007194245, 0.6934466361999512\n",
            "epoch 25, elapsed 0.61s: train = 0.8443308550185874, 0.5108898878097534, val = 0.7787769784172662, 0.6827135682106018\n",
            "epoch 26, elapsed 0.63s: train = 0.8494423791821561, 0.4901982545852661, val = 0.7841726618705036, 0.6715127229690552\n",
            "epoch 27, elapsed 0.66s: train = 0.8592007434944238, 0.46873584389686584, val = 0.7859712230215827, 0.6581224203109741\n",
            "epoch 28, elapsed 0.68s: train = 0.8643122676579925, 0.44640278816223145, val = 0.7823741007194245, 0.645258367061615\n",
            "epoch 29, elapsed 0.70s: train = 0.8712825278810409, 0.42530393600463867, val = 0.7823741007194245, 0.635916531085968\n",
            "epoch 30, elapsed 0.72s: train = 0.8801115241635687, 0.40585294365882874, val = 0.7841726618705036, 0.6300920248031616\n",
            "epoch 31, elapsed 0.74s: train = 0.8884758364312267, 0.3875647783279419, val = 0.789568345323741, 0.6257545948028564\n",
            "epoch 32, elapsed 0.76s: train = 0.8977695167286245, 0.36989328265190125, val = 0.7967625899280576, 0.6208022832870483\n",
            "epoch 33, elapsed 0.78s: train = 0.9010223048327137, 0.35272017121315, val = 0.8039568345323741, 0.6153444051742554\n",
            "epoch 34, elapsed 0.81s: train = 0.9056691449814126, 0.33582136034965515, val = 0.8111510791366906, 0.6096092462539673\n",
            "epoch 35, elapsed 0.83s: train = 0.912639405204461, 0.31959980726242065, val = 0.8129496402877698, 0.6057406067848206\n",
            "epoch 36, elapsed 0.85s: train = 0.9186802973977695, 0.3046039044857025, val = 0.8129496402877698, 0.604529857635498\n",
            "epoch 37, elapsed 0.87s: train = 0.9242565055762082, 0.2905968129634857, val = 0.8111510791366906, 0.6045737862586975\n",
            "epoch 38, elapsed 0.89s: train = 0.929368029739777, 0.27671560645103455, val = 0.8057553956834532, 0.6055150032043457\n",
            "epoch 39, elapsed 0.91s: train = 0.9340148698884758, 0.2628912329673767, val = 0.8057553956834532, 0.6054041385650635\n",
            "epoch 40, elapsed 0.93s: train = 0.9395910780669146, 0.2499629557132721, val = 0.8075539568345323, 0.604017436504364\n",
            "---Total training time: 0.9393796920776367 seconds ---\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!python3 apps/gn.py cuda gat"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8txKwHJvQlH-",
        "outputId": "2c88e8cd-32be-4d32-f515-c96f609a0640"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "---Selecting backend:  cuda() ---\n",
            "---Training GAT model---\n",
            "epoch 1, elapsed 0.13s: train = 0.0687732342007435, 2.616851806640625, val = 0.05575539568345324, 2.629812479019165\n",
            "epoch 2, elapsed 0.25s: train = 0.34572490706319703, 2.2221007347106934, val = 0.4064748201438849, 2.1131505966186523\n",
            "epoch 3, elapsed 0.36s: train = 0.27276951672862454, 2.2304086685180664, val = 0.22661870503597123, 2.3323240280151367\n",
            "epoch 4, elapsed 0.45s: train = 0.3805762081784387, 1.6058175563812256, val = 0.3489208633093525, 1.6444553136825562\n",
            "epoch 5, elapsed 0.55s: train = 0.6342936802973977, 1.3112070560455322, val = 0.6205035971223022, 1.3361002206802368\n",
            "epoch 6, elapsed 0.63s: train = 0.675185873605948, 1.1696418523788452, val = 0.64568345323741, 1.2157604694366455\n",
            "epoch 7, elapsed 0.72s: train = 0.7565055762081785, 1.011610507965088, val = 0.6996402877697842, 1.0769298076629639\n",
            "epoch 8, elapsed 0.80s: train = 0.7973977695167286, 0.887193500995636, val = 0.7392086330935251, 0.973892092704773\n",
            "epoch 9, elapsed 0.89s: train = 0.8257434944237918, 0.7635015249252319, val = 0.7931654676258992, 0.8690376281738281\n",
            "epoch 10, elapsed 0.97s: train = 0.8503717472118959, 0.6574841737747192, val = 0.8111510791366906, 0.7763491868972778\n",
            "epoch 11, elapsed 1.05s: train = 0.8605947955390335, 0.5858989357948303, val = 0.8219424460431655, 0.7215685248374939\n",
            "epoch 12, elapsed 1.14s: train = 0.862453531598513, 0.5367199182510376, val = 0.814748201438849, 0.6951377987861633\n",
            "epoch 13, elapsed 1.22s: train = 0.8675650557620818, 0.48237496614456177, val = 0.8057553956834532, 0.6664609909057617\n",
            "epoch 14, elapsed 1.31s: train = 0.8889405204460966, 0.4203212559223175, val = 0.8039568345323741, 0.6301099061965942\n",
            "epoch 15, elapsed 1.39s: train = 0.9154275092936803, 0.3616633713245392, val = 0.8165467625899281, 0.5890696048736572\n",
            "epoch 16, elapsed 1.47s: train = 0.9377323420074349, 0.31063827872276306, val = 0.8399280575539568, 0.5461532473564148\n",
            "epoch 17, elapsed 1.56s: train = 0.9488847583643123, 0.2688782811164856, val = 0.8489208633093526, 0.5082356929779053\n",
            "epoch 18, elapsed 1.64s: train = 0.9577137546468402, 0.2366422861814499, val = 0.8615107913669064, 0.4812597930431366\n",
            "epoch 19, elapsed 1.73s: train = 0.9595724907063197, 0.21191304922103882, val = 0.8579136690647482, 0.46571749448776245\n",
            "epoch 20, elapsed 1.81s: train = 0.9632899628252788, 0.19143977761268616, val = 0.8579136690647482, 0.45751598477363586\n",
            "epoch 21, elapsed 1.90s: train = 0.9688661710037175, 0.1702059954404831, val = 0.8633093525179856, 0.4503374993801117\n",
            "epoch 22, elapsed 1.98s: train = 0.9776951672862454, 0.14698998630046844, val = 0.8669064748201439, 0.44120508432388306\n",
            "epoch 23, elapsed 2.06s: train = 0.9851301115241635, 0.12574809789657593, val = 0.8705035971223022, 0.4329455494880676\n",
            "epoch 24, elapsed 2.14s: train = 0.9860594795539034, 0.10957199335098267, val = 0.8687050359712231, 0.429261714220047\n",
            "epoch 25, elapsed 2.23s: train = 0.9879182156133829, 0.09779305011034012, val = 0.8651079136690647, 0.43009617924690247\n",
            "epoch 26, elapsed 2.31s: train = 0.9879182156133829, 0.08857651799917221, val = 0.8669064748201439, 0.43432357907295227\n",
            "epoch 27, elapsed 2.39s: train = 0.9883828996282528, 0.08001706749200821, val = 0.8633093525179856, 0.43927836418151855\n",
            "epoch 28, elapsed 2.47s: train = 0.991635687732342, 0.07119034975767136, val = 0.8723021582733813, 0.44259169697761536\n",
            "epoch 29, elapsed 2.56s: train = 0.9930297397769516, 0.06267886608839035, val = 0.8723021582733813, 0.44425687193870544\n",
            "epoch 30, elapsed 2.64s: train = 0.9953531598513011, 0.05529332533478737, val = 0.8741007194244604, 0.44527485966682434\n",
            "epoch 31, elapsed 2.72s: train = 0.995817843866171, 0.04928426072001457, val = 0.8741007194244604, 0.4469199478626251\n",
            "epoch 32, elapsed 2.81s: train = 0.9967472118959108, 0.04446175694465637, val = 0.8723021582733813, 0.45019423961639404\n",
            "epoch 33, elapsed 2.90s: train = 0.9972118959107806, 0.040391914546489716, val = 0.8723021582733813, 0.4556218385696411\n",
            "epoch 34, elapsed 2.98s: train = 0.9976765799256505, 0.0367618165910244, val = 0.8723021582733813, 0.46374738216400146\n",
            "epoch 35, elapsed 3.07s: train = 0.9976765799256505, 0.03346538916230202, val = 0.8723021582733813, 0.4734887182712555\n",
            "epoch 36, elapsed 3.15s: train = 0.9976765799256505, 0.030428878962993622, val = 0.8723021582733813, 0.4833870232105255\n",
            "epoch 37, elapsed 3.23s: train = 0.9981412639405205, 0.027703961357474327, val = 0.8758992805755396, 0.4924549162387848\n",
            "epoch 38, elapsed 3.32s: train = 0.9986059479553904, 0.025321993976831436, val = 0.8687050359712231, 0.5005043148994446\n",
            "epoch 39, elapsed 3.40s: train = 0.9986059479553904, 0.023304428905248642, val = 0.8687050359712231, 0.5081102252006531\n",
            "epoch 40, elapsed 3.48s: train = 0.9986059479553904, 0.021593332290649414, val = 0.8705035971223022, 0.5158178210258484\n",
            "---Total training time: 3.5008535385131836 seconds ---\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "It can be seen that both Graph Convolution and Graph Attention can achieve higher validation accuracy than the baseline model. All the training times are within few seconds. The training time of Graph Convolution is not too much more than the baseline model. Graph Attention is slower, but can achieve the highest accuracy."
      ],
      "metadata": {
        "id": "gXDo8Y813KYi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Side Goal: Optimize Speed with Computation Libraries\n",
        "\n",
        "Note that the following results are produced on the personal computer of one of the team members (because we don't know how to install computation libraries on Google Colab). The CPU is Intel® Core™ i5-1135G7 @ 2.40GHz, which is not very powerful. There must be some number flucuations during the tests, but we believe the results are clear enough.\n",
        "\n",
        "Out initial plan was to optimize with sparse-computation-related libraries (such as cuSPARSE). However, we did some profiling after finishing the project, and the results show that dense operations were still dominant in time in graph networks.\n",
        "\n",
        "```\n",
        "> python -m cProfile -s tottime apps/gn.py cpu gc\n",
        "...\n",
        "         890088 function calls (867987 primitive calls) in 64.627 seconds\n",
        "\n",
        "   Ordered by: internal time\n",
        " \n",
        "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
        "      720   62.374    0.087   62.374    0.087 {built-in method needle.backend_ndarray.ndarray_backend_cpu.matmul}\n",
        "     1280    0.749    0.001    0.749    0.001 {built-in method needle.backend_ndarray.ndarray_backend_cpu.compact}\n",
        "     2708    0.335    0.000    0.335    0.000 data.py:379(<listcomp>)\n",
        "       40    0.138    0.003    0.139    0.003 {built-in method gc.collect}\n",
        "      240    0.075    0.000    0.075    0.000 {built-in method needle.backend_ndarray.ndarray_backend_cpu.sparse_dense_matmul}\n",
        "      120    0.070    0.001    0.070    0.001 {built-in method numpy.array}\n",
        "      800    0.060    0.000    0.060    0.000 {built-in method needle.backend_ndarray.ndarray_backend_cpu.scalar_power}\n",
        "     2360    0.057    0.000    0.057    0.000 {built-in method needle.backend_ndarray.ndarray_backend_cpu.ewise_add}\n",
        "...\n",
        "```\n",
        "\n",
        "The reader may refer to the documentation of `cProfile` at https://docs.python.org/3/library/profile.html. We care about `tottime`, which means the total time spent on the function itself (not including other functions it calls).\n",
        "\n",
        "It is clear that dense-dense matmul is the bottleneck. Based on the profiling results, we found that despite our initial plan to optimize with sparse-computation-related libraries, dense operations were still the dominant contributor to the computational time in our graph networks. This is not unexpected, as the number of floating-point operations (FLOPS) required for dense operations is higher than for sparse operations according to their definitions.\n",
        "\n",
        "As a result, we chose to optimize dense-dense matmul with Intel MKL, a highly optimized computing library for Intel CPU. This is exposed as a `matmul_mkl` function in our CPU backend. This piece of code is only compiled when the compiler detects `mkl.h` available. The new profiling result with the optimization is:\n",
        "\n",
        "```\n",
        "> python -m cProfile -s tottime apps/gn.py cpu gc\n",
        "...\n",
        "         882248 function calls (860147 primitive calls) in 3.071 seconds\n",
        "\n",
        "   Ordered by: internal time\n",
        "\n",
        "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
        "     1280    0.746    0.001    0.746    0.001 {built-in method needle.backend_ndarray.ndarray_backend_cpu.compact}\n",
        "      720    0.740    0.001    0.740    0.001 {built-in method needle.backend_ndarray.ndarray_backend_cpu.matmul_mkl}\n",
        "     2708    0.332    0.000    0.332    0.000 data.py:379(<listcomp>)\n",
        "       40    0.151    0.004    0.151    0.004 {built-in method gc.collect}\n",
        "      240    0.085    0.000    0.085    0.000 {built-in method needle.backend_ndarray.ndarray_backend_cpu.sparse_dense_matmul}\n",
        "      120    0.070    0.001    0.070    0.001 {built-in method numpy.array}\n",
        "      800    0.064    0.000    0.064    0.000 {built-in method needle.backend_ndarray.ndarray_backend_cpu.scalar_power}\n",
        "     2360    0.064    0.000    0.064    0.000 {built-in method needle.backend_ndarray.ndarray_backend_cpu.ewise_add}\n",
        "...\n",
        "```\n",
        "\n",
        "It is much (about 20x) faster, and the dense operations are still the bottleneck. So we gave up the idea of optimizing with sparse-computation-related libraries. We did similar profiling for the CUDA backend and GAT, which demonstrated similar results."
      ],
      "metadata": {
        "id": "k56D5emjmRGh"
      }
    }
  ]
}