02_2_Link_Prediction_GCNConv_GAE.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "02.2: Link_Prediction_.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.5"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/gulabpatel/Graph_Neural_Network/blob/main/02.2%3A%20Link_Prediction_GCNConv_GAE.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "75Nkk1S4QCCw"
      },
      "source": [
        "Video walkthrough the code : https://www.youtube.com/watch?v=m1G7oS9hmwE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C9d1rhsGfsnk",
        "outputId": "04071c54-7d98-45f3-d3e0-ed2e1e14ea30"
      },
      "source": [
        "# Install required packages.\n",
        "!pip install -q torch-scatter -f https://pytorch-geometric.com/whl/torch-1.9.0+cu102.html\n",
        "!pip install -q torch-sparse -f https://pytorch-geometric.com/whl/torch-1.9.0+cu102.html\n",
        "!pip install -q torch-geometric"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[K     |████████████████████████████████| 3.0 MB 2.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.6 MB 2.7 MB/s \n",
            "\u001b[K     |████████████████████████████████| 222 kB 5.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 376 kB 52.3 MB/s \n",
            "\u001b[K     |████████████████████████████████| 45 kB 3.0 MB/s \n",
            "\u001b[?25h  Building wheel for torch-geometric (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "32Ps2zmoJ9F0"
      },
      "source": [
        "import os.path as osp\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from sklearn.metrics import roc_auc_score\n",
        "\n",
        "from torch_geometric.utils import negative_sampling\n",
        "from torch_geometric.datasets import Planetoid\n",
        "import torch_geometric.transforms as T\n",
        "from torch_geometric.nn import GCNConv\n",
        "from torch_geometric.utils import train_test_split_edges"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jayw_l1IJ9F5"
      },
      "source": [
        "# GAE for link prediction\n",
        "\n",
        "[code](https://github.com/rusty1s/pytorch_geometric/blob/master/examples/link_pred.py)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BT30wjlqJ9F6"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device = \"cpu\""
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iQr_EyGlJ9F8",
        "outputId": "e95eddf9-e1f9-4bed-9cc1-2302f30c5c15"
      },
      "source": [
        "# load the Cora dataset\n",
        "dataset = 'Cora'\n",
        "path = osp.join('.', 'data', dataset)\n",
        "dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())\n",
        "data = dataset[0]\n",
        "print(dataset.data)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.x\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.tx\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.allx\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.y\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ty\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ally\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.graph\n",
            "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.test.index\n",
            "Processing...\n",
            "Done!\n",
            "Data(edge_index=[2, 10556], test_mask=[2708], train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EQJJ-3AyJ9F_",
        "outputId": "52acfd87-5eec-4473-c962-cf20b89610cb"
      },
      "source": [
        "# use train_test_split_edges to create neg and positive edges\n",
        "data.train_mask = data.val_mask = data.test_mask = data.y = None\n",
        "data = train_test_split_edges(data)\n",
        "print(data)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Data(test_neg_edge_index=[2, 527], test_pos_edge_index=[2, 527], train_neg_adj_mask=[2708, 2708], train_pos_edge_index=[2, 8976], val_neg_edge_index=[2, 263], val_pos_edge_index=[2, 263], x=[2708, 1433])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfDw7juZJ9GB"
      },
      "source": [
        "#### Simple autoencoder model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6K9dlFrTJ9GC"
      },
      "source": [
        "class Net(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = GCNConv(dataset.num_features, 128)\n",
        "        self.conv2 = GCNConv(128, 64)\n",
        "\n",
        "    def encode(self):\n",
        "        x = self.conv1(data.x, data.train_pos_edge_index) # convolution 1\n",
        "        x = x.relu()\n",
        "        return self.conv2(x, data.train_pos_edge_index) # convolution 2\n",
        "\n",
        "    def decode(self, z, pos_edge_index, neg_edge_index): # only pos and neg edges\n",
        "        edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1) # concatenate pos and neg edges\n",
        "        logits = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)  # dot product \n",
        "        return logits\n",
        "\n",
        "    def decode_all(self, z): \n",
        "        prob_adj = z @ z.t() # get adj NxN\n",
        "        return (prob_adj > 0).nonzero(as_tuple=False).t() # get predicted edge_list "
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NTa4sIN3J9GE"
      },
      "source": [
        "model, data = Net().to(device), data.to(device)\n",
        "optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZTSAbmEoJ9GG"
      },
      "source": [
        "def get_link_labels(pos_edge_index, neg_edge_index):\n",
        "    # returns a tensor:\n",
        "    # [1,1,1,1,...,0,0,0,0,0,..] with the number of ones is equel to the lenght of pos_edge_index\n",
        "    # and the number of zeros is equal to the length of neg_edge_index\n",
        "    E = pos_edge_index.size(1) + neg_edge_index.size(1)\n",
        "    link_labels = torch.zeros(E, dtype=torch.float, device=device)\n",
        "    link_labels[:pos_edge_index.size(1)] = 1.\n",
        "    return link_labels\n",
        "\n",
        "\n",
        "def train():\n",
        "    model.train()\n",
        "\n",
        "    neg_edge_index = negative_sampling(\n",
        "        edge_index=data.train_pos_edge_index, #positive edges\n",
        "        num_nodes=data.num_nodes, # number of nodes\n",
        "        num_neg_samples=data.train_pos_edge_index.size(1)) # number of neg_sample equal to number of pos_edges\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    z = model.encode() #encode\n",
        "    link_logits = model.decode(z, data.train_pos_edge_index, neg_edge_index) # decode\n",
        "    \n",
        "    link_labels = get_link_labels(data.train_pos_edge_index, neg_edge_index)\n",
        "    loss = F.binary_cross_entropy_with_logits(link_logits, link_labels)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    return loss\n",
        "\n",
        "\n",
        "@torch.no_grad()\n",
        "def test():\n",
        "    model.eval()\n",
        "    perfs = []\n",
        "    for prefix in [\"val\", \"test\"]:\n",
        "        pos_edge_index = data[f'{prefix}_pos_edge_index']\n",
        "        neg_edge_index = data[f'{prefix}_neg_edge_index']\n",
        "\n",
        "        z = model.encode() # encode train\n",
        "        link_logits = model.decode(z, pos_edge_index, neg_edge_index) # decode test or val\n",
        "        link_probs = link_logits.sigmoid() # apply sigmoid\n",
        "        \n",
        "        link_labels = get_link_labels(pos_edge_index, neg_edge_index) # get link\n",
        "        \n",
        "        perfs.append(roc_auc_score(link_labels.cpu(), link_probs.cpu())) #compute roc_auc score\n",
        "    return perfs\n"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BAdSegGVJ9GI",
        "outputId": "31474e08-74a3-412b-cad9-73c5f6b6ae74"
      },
      "source": [
        "best_val_perf = test_perf = 0\n",
        "for epoch in range(1, 101):\n",
        "    train_loss = train()\n",
        "    val_perf, tmp_test_perf = test()\n",
        "    if val_perf > best_val_perf:\n",
        "        best_val_perf = val_perf\n",
        "        test_perf = tmp_test_perf\n",
        "\n",
        "    if epoch % 10 == 0:\n",
        "        print('Epoch: {:03d}, Loss: {:.4f}, Val: {:.4f}, Test: {:.4f}'.format(epoch, train_loss, best_val_perf, test_perf))"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.\n",
            "To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /pytorch/aten/src/ATen/native/BinaryOps.cpp:467.)\n",
            "  return torch.floor_divide(self, other)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch: 010, Loss: 0.6899, Val: 0.7732, Test: 0.7663\n",
            "Epoch: 020, Loss: 0.6476, Val: 0.7732, Test: 0.7663\n",
            "Epoch: 030, Loss: 0.5748, Val: 0.7859, Test: 0.7619\n",
            "Epoch: 040, Loss: 0.5412, Val: 0.8017, Test: 0.7832\n",
            "Epoch: 050, Loss: 0.5077, Val: 0.8445, Test: 0.8210\n",
            "Epoch: 060, Loss: 0.4936, Val: 0.8469, Test: 0.8269\n",
            "Epoch: 070, Loss: 0.4762, Val: 0.8601, Test: 0.8576\n",
            "Epoch: 080, Loss: 0.4565, Val: 0.8773, Test: 0.8953\n",
            "Epoch: 090, Loss: 0.4424, Val: 0.8959, Test: 0.9039\n",
            "Epoch: 100, Loss: 0.4417, Val: 0.9026, Test: 0.9119\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wIRwHt8HET0-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "eec88f65-e730-4aa8-f983-3624358d2375"
      },
      "source": [
        "z = model.encode()\n",
        "final_edge_index = model.decode_all(z)\n",
        "final_edge_index"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[   0,    0,    0,  ..., 2707, 2707, 2707],\n",
              "        [   0,    2,    4,  ..., 2704, 2706, 2707]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    }
  ]
}