GNN/10_GAE_for_link_prediction.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9m0QemzX2XA2"
      },
      "outputs": [],
      "source": [
        "!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7vxlbmfd2XA2"
      },
      "outputs": [],
      "source": [
        "import os.path as osp\n",
        "\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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQTzAOVu2XA3"
      },
      "source": [
        "# GAE for link prediction\n",
        "\n",
        "[code](https://github.com/rusty1s/pytorch_geometric/blob/master/examples/link_pred.py)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OI6difUm2XA3"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device = \"cpu\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AAzWlIyI2XA4"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nT4dJ4uk2XA4"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xGZQZaYl2XA4"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crhBm-s32XA4"
      },
      "source": [
        "#### Simple autoencoder model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SdJNIZMY2XA4"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EU3H3Pud2XA4"
      },
      "outputs": [],
      "source": [
        "\n",
        "model, data = Net().to(device), data.to(device)\n",
        "optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lcXArTeC2XA4"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DrQv8RAL2XA4"
      },
      "outputs": [],
      "source": [
        "\n",
        "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E71ymOjM2XA5"
      },
      "outputs": [],
      "source": [
        "\n",
        "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",
        "    log = 'Epoch: {:03d}, Loss: {:.4f}, Val: {:.4f}, Test: {:.4f}'\n",
        "    if epoch % 10 == 0:\n",
        "        print(log.format(epoch, train_loss, best_val_perf, test_perf))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a_OZpmMv2XA5"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lyfNcdQS2XA5"
      },
      "outputs": [],
      "source": [
        "z = model.encode()\n",
        "final_edge_index = model.decode_all(z)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q7rRWarF2XA5"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qez05wPP2XA5"
      },
      "outputs": [],
      "source": []
    }
  ],
  "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.9.7"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}