GAT_using_dgl.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GAT using dgl.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyP5j/gMHs1hnMpf4YZVK3vy"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "uIDMOanZlDKN",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 217
        },
        "outputId": "96a939df-2c6c-43f3-a7c1-8848887e78c4"
      },
      "source": [
        "!pip install dgl\n",
        "import dgl\n",
        "from dgl.nn.pytorch import GATConv"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: dgl in /usr/local/lib/python3.6/dist-packages (0.4.3.post2)\n",
            "Requirement already satisfied: numpy>=1.14.0 in /usr/local/lib/python3.6/dist-packages (from dgl) (1.18.4)\n",
            "Requirement already satisfied: scipy>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from dgl) (1.4.1)\n",
            "Requirement already satisfied: networkx>=2.1 in /usr/local/lib/python3.6/dist-packages (from dgl) (2.4)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.6/dist-packages (from dgl) (2.23.0)\n",
            "Requirement already satisfied: decorator>=4.3.0 in /usr/local/lib/python3.6/dist-packages (from networkx>=2.1->dgl) (4.4.2)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->dgl) (2.9)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->dgl) (2020.4.5.1)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->dgl) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->dgl) (3.0.4)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H2OYxoDnlGCG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "class GATLayer(nn.Module):\n",
        "    def __init__(self, g, in_dim, out_dim):\n",
        "        super(GATLayer, self).__init__()\n",
        "        self.g = g\n",
        "        self.fc = nn.Linear(in_dim, out_dim, bias=False)\n",
        "        self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        \"\"\"Reinitialize learnable parameters.\"\"\"\n",
        "        gain = nn.init.calculate_gain('relu')\n",
        "        nn.init.xavier_normal_(self.fc.weight, gain=gain)\n",
        "        nn.init.xavier_normal_(self.attn_fc.weight, gain=gain)\n",
        "\n",
        "    def edge_attention(self, edges):\n",
        "        z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)\n",
        "        a = self.attn_fc(z2)\n",
        "        return {'e': F.leaky_relu(a)}\n",
        "\n",
        "    def message_func(self, edges):\n",
        "        return {'z': edges.src['z'], 'e': edges.data['e']}\n",
        "\n",
        "    def reduce_func(self, nodes):\n",
        "\n",
        "        alpha = F.softmax(nodes.mailbox['e'], dim=1)\n",
        "        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)\n",
        "        return {'h': h}\n",
        "\n",
        "    def forward(self, h):\n",
        "        z = self.fc(h)\n",
        "        self.g.ndata['z'] = z\n",
        "        self.g.apply_edges(self.edge_attention)\n",
        "        self.g.update_all(self.message_func, self.reduce_func)\n",
        "        return self.g.ndata.pop('h')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WYNgCraSyKLd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def edge_attention(self, edges):\n",
        "    z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)\n",
        "    a = self.attn_fc(z2)\n",
        "    return {'e' : F.leaky_relu(a)}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3JDJuyTiyKqi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def reduce_func(self, nodes):\n",
        " \n",
        "    alpha = F.softmax(nodes.mailbox['e'], dim=1)\n",
        "    h = torch.sum(alpha * nodes.mailbox['z'], dim=1)\n",
        "    return {'h' : h}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vGa2quFAyMM7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MultiHeadGATLayer(nn.Module):\n",
        "    def __init__(self, g, in_dim, out_dim, num_heads, merge='cat'):\n",
        "        super(MultiHeadGATLayer, self).__init__()\n",
        "        self.heads = nn.ModuleList()\n",
        "        for i in range(num_heads):\n",
        "            self.heads.append(GATLayer(g, in_dim, out_dim))\n",
        "        self.merge = merge\n",
        "\n",
        "    def forward(self, h):\n",
        "        head_outs = [attn_head(h) for attn_head in self.heads]\n",
        "        if self.merge == 'cat':\n",
        "            return torch.cat(head_outs, dim=1)\n",
        "        else:\n",
        "            return torch.mean(torch.stack(head_outs))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zv7jqZqGyNos",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class GAT(nn.Module):\n",
        "    def __init__(self, g, in_dim, hidden_dim, out_dim, num_heads):\n",
        "        super(GAT, self).__init__()\n",
        "        self.layer1 = MultiHeadGATLayer(g, in_dim, hidden_dim, num_heads)\n",
        "        self.layer2 = MultiHeadGATLayer(g, hidden_dim * num_heads, out_dim, 1)\n",
        "\n",
        "    def forward(self, h):\n",
        "        h = self.layer1(h)\n",
        "        h = F.elu(h)\n",
        "        h = self.layer2(h)\n",
        "        return h"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yGAeSslOyQA1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from dgl import DGLGraph\n",
        "from dgl.data import citation_graph as citegrh\n",
        "import networkx as nx\n",
        "\n",
        "def load_cora_data():\n",
        "    data = citegrh.load_cora()\n",
        "    features = torch.FloatTensor(data.features)\n",
        "    labels = torch.LongTensor(data.labels)\n",
        "    mask = torch.BoolTensor(data.train_mask)\n",
        "    g = DGLGraph(data.graph)\n",
        "    return g, features, labels, mask"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NgJBm2BoyQnl",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 685
        },
        "outputId": "b7884558-2886-45e7-90bb-3fdea9d6206c"
      },
      "source": [
        "import time\n",
        "import numpy as np\n",
        "\n",
        "g, features, labels, mask = load_cora_data()\n",
        "\n",
        "net = GAT(g,\n",
        "          in_dim=features.size()[1],\n",
        "          hidden_dim=8,\n",
        "          out_dim=7,\n",
        "          num_heads=2)\n",
        "\n",
        "optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)\n",
        "\n",
        "dur = []\n",
        "for epoch in range(30):\n",
        "    if epoch >= 3:\n",
        "        t0 = time.time()\n",
        "\n",
        "    logits = net(features)\n",
        "    logp = F.log_softmax(logits, 1)\n",
        "    loss = F.nll_loss(logp[mask], labels[mask])\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if epoch >= 3:\n",
        "        dur.append(time.time() - t0)\n",
        "\n",
        "    print(\"Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f}\".format(\n",
        "        epoch, loss.item(), np.mean(dur)))"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading /root/.dgl/cora.zip from https://data.dgl.ai/dataset/cora_raw.zip...\n",
            "Extracting file to /root/.dgl/cora\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py:3335: RuntimeWarning: Mean of empty slice.\n",
            "  out=out, **kwargs)\n",
            "/usr/local/lib/python3.6/dist-packages/numpy/core/_methods.py:161: RuntimeWarning: invalid value encountered in double_scalars\n",
            "  ret = ret.dtype.type(ret / rcount)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 00000 | Loss 1.9473 | Time(s) nan\n",
            "Epoch 00001 | Loss 1.9451 | Time(s) nan\n",
            "Epoch 00002 | Loss 1.9430 | Time(s) nan\n",
            "Epoch 00003 | Loss 1.9409 | Time(s) 0.1095\n",
            "Epoch 00004 | Loss 1.9388 | Time(s) 0.1053\n",
            "Epoch 00005 | Loss 1.9367 | Time(s) 0.1029\n",
            "Epoch 00006 | Loss 1.9345 | Time(s) 0.1018\n",
            "Epoch 00007 | Loss 1.9324 | Time(s) 0.1020\n",
            "Epoch 00008 | Loss 1.9303 | Time(s) 0.1022\n",
            "Epoch 00009 | Loss 1.9282 | Time(s) 0.1022\n",
            "Epoch 00010 | Loss 1.9260 | Time(s) 0.1018\n",
            "Epoch 00011 | Loss 1.9239 | Time(s) 0.1013\n",
            "Epoch 00012 | Loss 1.9217 | Time(s) 0.1014\n",
            "Epoch 00013 | Loss 1.9196 | Time(s) 0.1027\n",
            "Epoch 00014 | Loss 1.9174 | Time(s) 0.1024\n",
            "Epoch 00015 | Loss 1.9153 | Time(s) 0.1022\n",
            "Epoch 00016 | Loss 1.9131 | Time(s) 0.1018\n",
            "Epoch 00017 | Loss 1.9109 | Time(s) 0.1018\n",
            "Epoch 00018 | Loss 1.9087 | Time(s) 0.1017\n",
            "Epoch 00019 | Loss 1.9065 | Time(s) 0.1014\n",
            "Epoch 00020 | Loss 1.9043 | Time(s) 0.1014\n",
            "Epoch 00021 | Loss 1.9021 | Time(s) 0.1011\n",
            "Epoch 00022 | Loss 1.8998 | Time(s) 0.1010\n",
            "Epoch 00023 | Loss 1.8976 | Time(s) 0.1014\n",
            "Epoch 00024 | Loss 1.8953 | Time(s) 0.1012\n",
            "Epoch 00025 | Loss 1.8931 | Time(s) 0.1010\n",
            "Epoch 00026 | Loss 1.8908 | Time(s) 0.1011\n",
            "Epoch 00027 | Loss 1.8885 | Time(s) 0.1010\n",
            "Epoch 00028 | Loss 1.8862 | Time(s) 0.1009\n",
            "Epoch 00029 | Loss 1.8838 | Time(s) 0.1010\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}