labs/lab3_node_class.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Lab 3 - Node Classification\n",
    "\n",
    "12/14/2022\n",
    "\n",
    "Jay Urbain, PhD\n",
    "\n",
    "This dataset was first introduced by [Yang et al. (2016)](https://arxiv.org/abs/1603.08861) as one of the datasets of the `Planetoid` benchmark suite.\n",
    "\n",
    "Complete all ## TODO ## sections below.\n",
    "\n",
    "References:\n",
    "\n",
    "https://pytorch.org\n",
    "\n",
    "https://pytorch-geometric.readthedocs.io/en/latest/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F1op-CbyLuN4",
    "outputId": "77aaf048-dc7e-491d-a0ed-7a13f3692cc0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.12.0\n"
     ]
    }
   ],
   "source": [
    "# Install required packages.\n",
    "import os\n",
    "import torch\n",
    "os.environ['TORCH'] = torch.__version__\n",
    "print(torch.__version__)\n",
    "\n",
    "# !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
    "# !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
    "# !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git\n",
    "\n",
    "# Helper function for visualization.\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "def visualize(h, color):\n",
    "    z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy())\n",
    "\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "\n",
    "    plt.scatter(z[:, 0], z[:, 1], s=70, c=color, cmap=\"Set2\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dszt2RUHE7lW"
   },
   "source": [
    "# Node Classification with Graph Neural Networks\n",
    "\n",
    "Apply **Graph Neural Networks (GNNs) to the task of node classification**.\n",
    "Here, we are given the ground-truth labels of only a small subset of nodes, and want to infer the labels for all the remaining nodes (*transductive learning*).\n",
    "\n",
    "To demonstrate, we make use of the `Cora` dataset, which is a **citation network** where nodes represent documents.\n",
    "Each node is described by a 1433-dimensional bag-of-words feature vector.\n",
    "Two documents are connected if there exists a citation link between them.\n",
    "The task is to infer the category of each document (7 in total)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "imGrKO5YH11-",
    "outputId": "2a64e675-158c-4519-ab28-1923c4ab830a"
   },
   "outputs": [],
   "source": [
    "from torch_geometric.datasets import Planetoid\n",
    "from torch_geometric.transforms import NormalizeFeatures\n",
    "\n",
    "dataset = Planetoid(root='data/Planetoid', name='Cora', transform=NormalizeFeatures())\n",
    "\n",
    "## TODO\n",
    "# Gather properties of the graph - fill in the formatted print statements\n",
    "print()\n",
    "print(f'Dataset: {}:')\n",
    "print('======================')\n",
    "print(f'Number of graphs: {}')\n",
    "print(f'Number of features: {}')\n",
    "print(f'Number of classes: {}')\n",
    "\n",
    "data = dataset[0]  # Get the first graph object.\n",
    "\n",
    "print()\n",
    "print(data)\n",
    "print('===========================================================================================================')\n",
    "\n",
    "## TODO\n",
    "# Gather some statistics about the graph - fill in the formatted print statements\n",
    "print(f'Number of nodes: {}')\n",
    "print(f'Number of edges: {}')\n",
    "print(f'Average node degree: {}')\n",
    "print(f'Number of training nodes: {}')\n",
    "print(f'Training node label rate: {}')\n",
    "print(f'Has isolated nodes: {}')\n",
    "print(f'Has self-loops: {}')\n",
    "print(f'Is undirected: {}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eqWR0j_kIx67"
   },
   "source": [
    "Overall, this dataset is quite similar to the  [`KarateClub`](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.KarateClub) network.\n",
    "The `Cora` network holds 2,708 nodes and 10,556 edges, resulting in an average node degree of 3.9.\n",
    "For training this dataset, we are given the ground-truth categories of 140 nodes (20 for each class).\n",
    "This results in a training node label rate of only 5%.\n",
    "\n",
    "In contrast to `KarateClub`, this graph holds the additional attributes `val_mask` and `test_mask`, which denotes which nodes should be used for validation and testing.\n",
    "\n",
    "We make use of **[data transformations](https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html#data-transforms) via `transform=NormalizeFeatures()`**.\n",
    "Transforms can be used to modify your input data before inputting them into a neural network, *e.g.*, for normalization or data augmentation.\n",
    "The [row-normalize](https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.NormalizeFeatures) the bag-of-words input feature vectors.\n",
    "\n",
    "We can further see that this network is undirected, and that there exists no isolated nodes (each document has at least one citation)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5IRdAELVKOl6"
   },
   "source": [
    "## Training a Multi-layer Perception Network (MLP)\n",
    "\n",
    "In theory, we should be able to infer the category of a document solely based on its content, *i.e.* its bag-of-words feature representation, without taking any relational information into account.\n",
    "\n",
    "Let's verify that by constructing a simple MLP that solely operates on input node features (using shared weights across all nodes):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "afXwPCA3KNoC",
    "outputId": "d6dafb2a-1415-48b1-cfa2-7e0baf7f8491"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.nn import Linear\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class MLP(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super().__init__()\n",
    "        torch.manual_seed(12345)\n",
    "        self.lin1 = Linear(dataset.num_features, hidden_channels)\n",
    "        self.lin2 = Linear(hidden_channels, dataset.num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.lin1(x)\n",
    "        x = x.relu()\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin2(x)\n",
    "        return x\n",
    "\n",
    "model = MLP(hidden_channels=16)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L_PO9EEHL7J6"
   },
   "source": [
    "Our MLP is defined by two linear layers and enhanced by [ReLU](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html?highlight=relu#torch.nn.ReLU) non-linearity and [dropout](https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html?highlight=dropout#torch.nn.Dropout).\n",
    "Here, we first reduce the 1433-dimensional feature vector to a low-dimensional embedding (`hidden_channels=16`), while the second linear layer acts as a classifier that should map each low-dimensional node embedding to one of the 7 classes.\n",
    "\n",
    "Let's train our simple MLP.\n",
    "\n",
    "We make use of the **cross entropy loss** and **Adam optimizer**.\n",
    "\n",
    "Define a **`test` function** to evaluate how well our final model performs on the test node set (which labels have not been observed during training)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 300
    },
    "id": "0YgHcLXMLk4o",
    "outputId": "71316c1f-addc-4aab-c61e-558c4ca139e5"
   },
   "outputs": [],
   "source": [
    "from IPython.display import Javascript  # Restrict height of output cell.\n",
    "\n",
    "model = MLP(hidden_channels=16)\n",
    "criterion = torch.nn.CrossEntropyLoss()  # Define loss criterion.\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)  # Define optimizer.\n",
    "\n",
    "def train():\n",
    "      model.train()\n",
    "      optimizer.zero_grad()  # Clear gradients.\n",
    "      out = model(data.x)  # Perform a single forward pass.\n",
    "      loss = criterion(out[data.train_mask], data.y[data.train_mask])  # Compute the loss solely based on the training nodes.\n",
    "      loss.backward()  # Derive gradients.\n",
    "      optimizer.step()  # Update parameters based on gradients.\n",
    "      return loss\n",
    "\n",
    "\n",
    "def test():\n",
    "      ## TODO ##\n",
    "      # complete the test function\n",
    "      model.eval()\n",
    "      # your code here\n",
    "    \n",
    "        \n",
    "      return test_acc\n",
    "\n",
    "for epoch in range(1, 201):\n",
    "    loss = train()\n",
    "    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kG4IKy9YOLGF"
   },
   "source": [
    "After training the model, we can call the `test` function to see how well our model performs on unseen labels.\n",
    "Here, we are interested in the accuracy of the model, *i.e.*, the ratio of correctly classified nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dBBCeLlAL0oL",
    "outputId": "82b69943-feb5-4df4-fef0-5178682ba47e"
   },
   "outputs": [],
   "source": [
    "test_acc = test()\n",
    "print(f'Test Accuracy: {test_acc:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_jjJOB-VO-cw"
   },
   "source": [
    "As one can see, our MLP performs rather bad with only about 59% test accuracy.\n",
    "But why does the MLP do not perform better?\n",
    "The main reason for that is that this model suffers from heavy overfitting due to only having access to a **small amount of training nodes**, and therefore generalizes poorly to unseen node representations.\n",
    "\n",
    "It also fails to incorporate an important bias into the model: **Cited papers are very likely related to the category of a document**.\n",
    "That is exactly where Graph Neural Networks come into play and can help to boost the performance of our model.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_OWGw54wRd98"
   },
   "source": [
    "## Training a Graph Neural Network (GNN)\n",
    "\n",
    "We can easily convert our MLP to a GNN by swapping the `torch.nn.Linear` layers with PyG's GNN operators.\n",
    "\n",
    "Replace the linear layers by the [`GCNConv`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GCNConv) module.\n",
    "To recap, the **GCN layer** ([Kipf et al. (2017)](https://arxiv.org/abs/1609.02907)) is defined as\n",
    "\n",
    "$$\n",
    "\\mathbf{x}_v^{(\\ell + 1)} = \\mathbf{W}^{(\\ell + 1)} \\sum_{w \\in \\mathcal{N}(v) \\, \\cup \\, \\{ v \\}} \\frac{1}{c_{w,v}} \\cdot \\mathbf{x}_w^{(\\ell)}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{W}^{(\\ell + 1)}$ denotes a trainable weight matrix of shape `[num_output_features, num_input_features]` and $c_{w,v}$ refers to a fixed normalization coefficient for each edge.\n",
    "In contrast, a single `Linear` layer is defined as\n",
    "\n",
    "$$\n",
    "\\mathbf{x}_v^{(\\ell + 1)} = \\mathbf{W}^{(\\ell + 1)} \\mathbf{x}_v^{(\\ell)}\n",
    "$$\n",
    "\n",
    "which does not make use of neighboring node information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fmXWs1dKIzD8",
    "outputId": "d3936b27-6d6b-4159-fb0d-08e0e381a715"
   },
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv\n",
    "\n",
    "\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super().__init__()\n",
    "        torch.manual_seed(1234567)\n",
    "        self.conv1 = GCNConv(dataset.num_features, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, dataset.num_classes)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "        return x\n",
    "\n",
    "model = GCN(hidden_channels=16)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XhO8QDgYf_Q8"
   },
   "source": [
    "Let's visualize the node embeddings of our **untrained** GCN network.\n",
    "For visualization, we make use of [**TSNE**](https://scikit-learn.org/stable/modules/generated/sklearn.manifold.TSNE.html) to embed our 7-dimensional node embeddings onto a 2D plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 672
    },
    "id": "ntt9qVFXlk6A",
    "outputId": "98abab94-b8c1-403a-b961-865a34a3b21b"
   },
   "outputs": [],
   "source": [
    "model = GCN(hidden_channels=16)\n",
    "model.eval()\n",
    "\n",
    "out = model(data.x, data.edge_index)\n",
    "visualize(out, color=data.y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fpdscco5g6kG"
   },
   "source": [
    "We  can do better by training our model.\n",
    "The training and testing procedure is once again the same, but this time we make use of the node features `x` **and** the graph connectivity `edge_index` as input to our GCN model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 300
    },
    "id": "p3TAi69zI1bO",
    "outputId": "4ad25e08-92dd-4396-f144-7a98126d6bf4"
   },
   "outputs": [],
   "source": [
    "from IPython.display import Javascript  # Restrict height of output cell.\n",
    "\n",
    "model = GCN(hidden_channels=16)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "def train():\n",
    "      ## TODO ##\n",
    "      # complete the training function\n",
    "      model.train()\n",
    "      ################\n",
    "      # your code here\n",
    "    \n",
    "        \n",
    "      ################\n",
    "      return loss\n",
    "\n",
    "def test():\n",
    "      model.eval()\n",
    "      out = model(data.x, data.edge_index)\n",
    "      pred = out.argmax(dim=1)  # Use the class with highest probability.\n",
    "      test_correct = pred[data.test_mask] == data.y[data.test_mask]  # Check against ground-truth labels.\n",
    "      test_acc = int(test_correct.sum()) / int(data.test_mask.sum())  # Derive ratio of correct predictions.\n",
    "      return test_acc\n",
    "\n",
    "\n",
    "for epoch in range(1, 101):\n",
    "    loss = train()\n",
    "    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "opBBGQHqg5ZO"
   },
   "source": [
    "After training the model, we can check its test accuracy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8zOh6IIeI3Op",
    "outputId": "1495b4b5-77c3-4374-ced1-5ad9d352a465"
   },
   "outputs": [],
   "source": [
    "test_acc = test()\n",
    "print(f'Test Accuracy: {test_acc:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yhofzjaqhfY2"
   },
   "source": [
    "By simply swapping the linear layers with GNN layers, we can reach **81.5% of test accuracy**!\n",
    "This is in stark contrast to the 59% of test accuracy obtained by our MLP, indicating that relational information plays a crucial role in obtaining better performance.\n",
    "\n",
    "We can also verify that once again by looking at the output embeddings of our **trained** model, which now produces a far better clustering of nodes of the same category."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 578
    },
    "id": "9r_VmGMukf5R",
    "outputId": "b4f41cab-6203-42df-d5a8-86c6b8f13cdd"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "out = model(data.x, data.edge_index)\n",
    "visualize(out, color=data.y)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}