SSL_for_GNNs_replicating_GraphMAE.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "this notebook is a simple replication of the [GraphMAE paper](https://arxiv.org/abs/2205.10803)\n",
        "\n",
        "which uses a masked GNN autoencoder \n",
        "\n",
        "its pretraining task is to reconstruct masked node features from neighboring nodes "
      ],
      "metadata": {
        "id": "mbTb_bxP2jG2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "before passing the graph (nodes and edge_index) to the encoder it masks some of the nodes (it replaces its features with a learnable vector like the  \\[MASK] token in BERT) "
      ],
      "metadata": {
        "id": "OE8oZewNIKSf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "the notebook is focused on the effectiveness of SSL so it might lack some improvements for the classifiers"
      ],
      "metadata": {
        "id": "q3FhnIroI7zg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "[colab link](https://colab.research.google.com/drive/1-uoWt0KtpAYs0egjAhJPhGGUFhHgxVwF?usp=sharing)"
      ],
      "metadata": {
        "id": "v8FjeCjrkSwc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torch-scatter torch-sparse torch-cluster "
      ],
      "metadata": {
        "id": "9x2hMmdevmFe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GPaBh7mmrWh2",
        "outputId": "ff74066b-6945-433f-b3c2-eaefcc97cb34"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "gpu up\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7f9b7cdc6550>"
            ]
          },
          "metadata": {},
          "execution_count": 110
        }
      ],
      "source": [
        "from sklearn.preprocessing import StandardScaler\n",
        "import torch \n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "import torch_geometric.nn as gnn\n",
        "from torch_geometric.datasets import Planetoid\n",
        "from torch import nn\n",
        "from typing import Optional,Tuple,List\n",
        "import torch_geometric\n",
        "\n",
        "if torch.cuda.is_available():  \n",
        "  dev = \"cuda:0\" \n",
        "  print(\"gpu up\")\n",
        "else:  \n",
        "  dev = \"cpu\"  \n",
        "\n",
        "# dev = \"cpu\"  \n",
        "\n",
        "device = torch.device(dev)\n",
        "\n",
        "\n",
        "scaler =  StandardScaler()\n",
        "\n",
        "\n",
        "\n",
        "torch.manual_seed(20)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jPHnSknTrWh5"
      },
      "outputs": [],
      "source": [
        "\n",
        "def sce(x,y,alpha=2):\n",
        "    # scaled cosine error \n",
        "    # mean( (1 - cosine similarity)^alpha ) \n",
        "    xnorm = torch.max(x.norm(p=2, dim=-1).unsqueeze(-1),torch.tensor(1e-12))\n",
        "    ynorm = torch.max(y.norm(p=2, dim=-1).unsqueeze(-1),torch.tensor(1e-12))\n",
        "    \n",
        "    cosine_error = 1- (x*y/(xnorm * ynorm)).sum(dim=-1)\n",
        "    return  cosine_error.pow(alpha).mean()\n",
        "\n",
        "#these .float() calls are just for unifying the tensor datatypes \n",
        "class GConv(nn.Module):\n",
        "\n",
        "    def __init__(self, emb_dim:int=500, num_layers:int=2, encode:bool=True, concat_out:bool=False, device='cpu', dropout=0.2):\n",
        "     \n",
        "        super(GConv,self).__init__()\n",
        "        self.num_layers = num_layers\n",
        "        self.gconv_layers = []\n",
        "        self.norm_layers = []\n",
        "        self.encode = encode\n",
        "        for _ in range(num_layers):\n",
        "            self.gconv_layers.append(gnn.SAGEConv(emb_dim, emb_dim, dropout=dropout,).to(device))\n",
        "            if self.encode:\n",
        "                self.norm_layers.append(nn.LayerNorm(emb_dim).to(device))\n",
        "\n",
        "        self.concat_out = concat_out\n",
        "    \n",
        "    def forward(self, x, edge_index):\n",
        "        \n",
        "        outs = []\n",
        "        if self.encode:\n",
        "            outs.append(self.norm_layers[0](self.gconv_layers[0](x, edge_index)))\n",
        "        else:\n",
        "            outs.append(self.gconv_layers[0](x, edge_index))\n",
        "        for i in range(1,self.num_layers):\n",
        "            if self.encode:\n",
        "                outs.append(self.norm_layers[i](self.gconv_layers[i](outs[-1], edge_index)))\n",
        "            else:\n",
        "                outs.append(self.gconv_layers[i](outs[-1], edge_index))\n",
        "        if self.concat_out:\n",
        "            return torch.cat(outs, dim = -1)\n",
        "        \n",
        "        return outs[-1]\n",
        "        \n",
        "class GraphMAE(nn.Module):\n",
        "\n",
        "    def __init__(self, \n",
        "            emb_dim:int=500,\n",
        "            masked_ratio:float=0.3, \n",
        "            num_encode_layers:int=2, \n",
        "            concat_out:bool=False, \n",
        "            device:torch.device='cpu',\n",
        "            dropout:float=0.2\n",
        "            ):\n",
        "        super(GraphMAE, self).__init__()\n",
        "        self.device = device\n",
        "        self.encoder = GConv(emb_dim, num_encode_layers, concat_out=concat_out, encode=True, device=device, dropout=dropout).float().to(device)\n",
        "        self.decoder = GConv(emb_dim, 1, encode=False, device=device, dropout=0).float().to(device)\n",
        "        self.masked_ratio = masked_ratio \n",
        "        self.encoder_mask_token = nn.Parameter(torch.zeros(1,emb_dim)).float().to(device) \n",
        "        self.remask_token = nn.Parameter(torch.zeros(1,emb_dim)).float().to(device)\n",
        "        if concat_out:\n",
        "            self.encoder_to_decoder = nn.Linear(num_encode_layers*emb_dim, emb_dim).to(device)\n",
        "        else:\n",
        "            self.encoder_to_decoder = nn.Linear(emb_dim, emb_dim).to(device)\n",
        "\n",
        "    \n",
        "    def forward(self, x, edge_index):\n",
        "        num_nodes = len(x)\n",
        "        masked_nodes_index = torch.randperm(num_nodes)[:int(num_nodes * self.masked_ratio)]\n",
        "        recon_x = x.clone().float().to(self.device)\n",
        "        recon_x[masked_nodes_index,:] = torch.tensor(0).float()\n",
        "        recon_x[masked_nodes_index,:] += self.encoder_mask_token\n",
        "        \n",
        "        recon_x = self.encoder(recon_x, edge_index)\n",
        "\n",
        "        recon_x = self.encoder_to_decoder(recon_x)\n",
        "\n",
        "        #remasking         \n",
        "        recon_x[masked_nodes_index,:] = 0\n",
        "        recon_x[masked_nodes_index,:] += self.remask_token\n",
        "\n",
        "        recon_x = self.decoder(recon_x, edge_index)\n",
        "\n",
        "        return recon_x, masked_nodes_index\n",
        "\n",
        "    def encode(self, x, edge_index):\n",
        "        return self.encoder_to_decoder(self.encoder(x.float(), edge_index.long()))\n",
        "    \n",
        "    def decode(self, x, edge_index):\n",
        "      return self.decode(x, edge_index)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5XVtbu0vrWh9"
      },
      "outputs": [],
      "source": [
        "dataset = Planetoid(\"./pubmed\",\"PubMed\")\n",
        "graph  =  dataset[0]\n",
        "\n",
        "graphmae = GraphMAE(masked_ratio=0.4, num_encode_layers=2, dropout=0.2, concat_out=True, device=device).float().to(device)\n",
        "scaler = StandardScaler() \n",
        "scaler.fit(graph.x)\n",
        "\n",
        "scaled_x = torch.tensor(scaler.transform(graph.x)).to(device)\n",
        "graph = graph.to(device)\n",
        "# edge_index = torch.tensor(graph.edge_index, dtype= torch.double)\n",
        "# some modules might use lazy initialization (which means parameters are initialized on the first forward call) \n",
        "graphmae(scaled_x, graph.edge_index.to(device))\n",
        "optimizer = torch.optim.Adam(graphmae.parameters(), lr=0.001, weight_decay=1e-5)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "epochs  = 301\n",
        "svc_accs = []\n",
        "lr_accs = []\n",
        "for e in range(epochs):\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    out, masked_nodes_index = graphmae(scaled_x, graph.edge_index.to(device))\n",
        "    # break\n",
        "    err = sce(out[masked_nodes_index,:], scaled_x[masked_nodes_index,:], alpha=2)    \n",
        "    err.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if e % 50 == 0: \n",
        "      print(\"scaled cosine error:\",err.item())\n",
        "      "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H7dzVTTFECTB",
        "outputId": "97bcb35f-af8e-4fdf-eaf8-0014964ecbb9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "scaled cosine error: 0.9978873160494309\n",
            "scaled cosine error: 0.6998318488349362\n",
            "scaled cosine error: 0.6720454481031665\n",
            "scaled cosine error: 0.6661558207945321\n",
            "scaled cosine error: 0.6524083907653202\n",
            "scaled cosine error: 0.6536281098493184\n",
            "scaled cosine error: 0.6462239074967576\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.linear_model import LogisticRegression\n",
        "from sklearn.tree import DecisionTreeClassifier\n",
        "from sklearn.svm import SVC\n",
        "from sklearn.metrics import accuracy_score\n",
        "\n",
        "def evaluate_embeddings(x, edge_index, target, train_mask, val_mask):\n",
        "  graphmae.eval()\n",
        "  torch.manual_seed(20)\n",
        "  encoded_embeddings = graphmae.encode(x, edge_index)\n",
        "  # print(encoded_embeddings[:5,:5])\n",
        "  lr = LogisticRegression()\n",
        "  svc = SVC()\n",
        "  lr.fit(encoded_embeddings[train_mask].detach().cpu().numpy(), target[train_mask].cpu().numpy())\n",
        "  preds = lr.predict(encoded_embeddings[val_mask].detach().cpu().numpy())\n",
        "  print(\"Logistic Regression Accuracy:\",accuracy_score(target[val_mask].cpu().numpy(), preds))\n",
        " \n",
        "  svc.fit(encoded_embeddings[train_mask].detach().cpu().numpy(), target[train_mask].cpu().numpy())\n",
        "  preds = svc.predict(encoded_embeddings[val_mask].detach().cpu().numpy())\n",
        "  print(\"SVC Accuracy:\",accuracy_score(target[val_mask].cpu().numpy(), preds))\n",
        "  "
      ],
      "metadata": {
        "id": "PkoYfs-ao-m9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def fit_classifiers(nodes, y, train_mask, val_mask):\n",
        "  svc = SVC()\n",
        "  lr = LogisticRegression()\n",
        "  print(\"scores on non-encoded vectors\")\n",
        "  lr.fit(nodes[train_mask].detach().cpu().numpy(), y[train_mask].cpu().numpy())\n",
        "  preds = lr.predict(nodes[val_mask].detach().cpu().numpy())\n",
        "  print(\"Logistic regression \",accuracy_score(y[val_mask].cpu().numpy(), preds))\n",
        "\n",
        "  svc.fit(nodes[train_mask].detach().cpu().numpy(), y[train_mask].cpu().numpy())\n",
        "  preds = svc.predict(nodes[val_mask].detach().cpu().numpy())\n",
        "  print(\"SVC score\",accuracy_score(y[val_mask].cpu().numpy(), preds))\n"
      ],
      "metadata": {
        "id": "21Ij3JbimrXd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "comparing generated embeddings to just using scaled features for classification"
      ],
      "metadata": {
        "id": "x8zi0T-lHX3Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "evaluate_embeddings(scaled_x, graph.edge_index.to(device), graph.y, graph.train_mask, graph.test_mask)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WumXDqoJEPYu",
        "outputId": "92b99d92-9f41-4aaf-850f-51b161f2bd5a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Logistic Regression Accuracy: 0.743\n",
            "SVC Accuracy: 0.76\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "fit_classifiers(scaled_x, graph.y, graph.train_mask, graph.test_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e5dQIaMbq13I",
        "outputId": "eb392eb8-7b87-4cd6-e8f6-91b4fff567eb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "scores on non-encoded vectors\n",
            "Logistic regression  0.702\n",
            "SVC score 0.677\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "but if increase training examples(by using val_mask) pretraining value is much less"
      ],
      "metadata": {
        "id": "IwXr8onElPEb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "evaluate_embeddings(scaled_x, graph.edge_index, graph.y , graph.val_mask, graph.test_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2-TL_8gZqUQ-",
        "outputId": "4d1127a1-4843-4cf0-8679-79d8a4887824"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Logistic Regression Accuracy: 0.737\n",
            "SVC Accuracy: 0.803\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "fit_classifiers(scaled_x, graph.y, graph.val_mask, graph.test_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mLfEiJi8lO1D",
        "outputId": "2e3c2f62-ac0f-4237-9721-0f6804ba0871"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "scores on non-encoded vectors\n",
            "Logistic regression  0.768\n",
            "SVC score 0.762\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Doing the same comparison using a GNN to for classification"
      ],
      "metadata": {
        "id": "-_Vbk6VlHtQH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class GNNClassifier(nn.Module):\n",
        "  def __init__(self, \n",
        "              emb_dim:int=500,\n",
        "              masked_ratio:float=0.3, \n",
        "              num_encode_layers:int=2, \n",
        "              concat_out:bool=False, \n",
        "              device:torch.device='cpu',\n",
        "              dropout:float=0.2\n",
        "              ):\n",
        "    super(GNNClassifier, self).__init__()\n",
        "    self.device = device\n",
        "    self.gnn_conv = GConv(emb_dim, num_encode_layers, concat_out=concat_out, encode=True, device=device, dropout=dropout).float().to(device)\n",
        "    self.criterion = nn.CrossEntropyLoss()\n",
        "    self.classifier = nn.Linear(emb_dim, 3).to(device)\n",
        "\n",
        "  def forward(self, x, edge_index):\n",
        "    return self.classifier(self.gnn_conv(x.float(), edge_index.long()))\n",
        "\n",
        "  def loss(self, x, edge_index, target, mask):\n",
        "\n",
        "    outs =  self.classifier(self.gnn_conv(x.float(), edge_index.long()))\n",
        "\n",
        "    return outs, self.criterion(outs[mask], target[mask])\n",
        "      \n"
      ],
      "metadata": {
        "id": "AIhTPKMDsUsl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "gnn_classifier = GNNClassifier(device=device)\n",
        "epochs  = 50  # low number of epochs to not give much of a chance to overfit\n",
        "svc_accs = []\n",
        "lr_accs = []\n",
        "optimizer = torch.optim.Adam(gnn_classifier.parameters(), lr=0.001, weight_decay=1e-5)\n",
        "for e in range(epochs):\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    outs, err = gnn_classifier.loss(scaled_x, graph.edge_index.to(device), graph.y, graph.train_mask)\n",
        "    # break\n",
        "    err.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if e % 10 == 0: \n",
        "      print(\"Cross Entropy:\",err.item())\n",
        "      print(accuracy_score(outs[graph.test_mask].argmax(-1).cpu().detach().numpy(),graph.y[graph.test_mask].cpu()))\n",
        "      "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Cn31OCRexE4_",
        "outputId": "2edbf073-fc5f-44f1-bbbe-19e5aeca5bde"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cross Entropy: 1.2189218997955322\n",
            "0.345\n",
            "Cross Entropy: 0.535369336605072\n",
            "0.536\n",
            "Cross Entropy: 0.21646060049533844\n",
            "0.596\n",
            "Cross Entropy: 0.08846046775579453\n",
            "0.61\n",
            "Cross Entropy: 0.03282538428902626\n",
            "0.616\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "gnn_classifier = GNNClassifier(device=device)\n",
        "epochs  = 50\n",
        "svc_accs = []\n",
        "lr_accs = []\n",
        "optimizer = torch.optim.Adam(gnn_classifier.parameters(), lr=0.001, weight_decay=1e-5)\n",
        "\n",
        "with torch.no_grad():\n",
        "  graphmae.eval()\n",
        "  ssl_embeddings = graphmae.encode(scaled_x, graph.edge_index.to(device))\n",
        "\n",
        "for e in range(epochs):\n",
        "    \n",
        "    optimizer.zero_grad()\n",
        "    outs, err = gnn_classifier.loss(ssl_embeddings, graph.edge_index.to(device), graph.y, graph.train_mask)\n",
        "    # break\n",
        "    err.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if e % 10 == 0: \n",
        "      print(\"Cross Entropy:\",err.item())\n",
        "      # print(accuracy_score(outs[graph.test_mask+graph.val_mask].argmax(-1).cpu().detach().numpy(),graph.y[graph.test_mask+graph.val_mask].cpu()))\n",
        "      print(accuracy_score(outs[graph.test_mask].argmax(-1).cpu().detach().numpy(),graph.y[graph.test_mask].cpu()))\n",
        "      "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FW5UrLrsdj2V",
        "outputId": "180de918-b8bd-4fd4-cc4b-d10ee2d7bdba"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cross Entropy: 1.1787744760513306\n",
            "0.476\n",
            "Cross Entropy: 0.469384104013443\n",
            "0.722\n",
            "Cross Entropy: 0.36979395151138306\n",
            "0.714\n",
            "Cross Entropy: 0.3411937654018402\n",
            "0.726\n",
            "Cross Entropy: 0.26469433307647705\n",
            "0.738\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.8.13 ('gnns')",
      "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.13"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "65071604a904ffcbac85598bd4b7cd132c32dd5e96ad1769b5378a4247c1d236"
      }
    },
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}