CS224W_Colab_0.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "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.7.3"
    },
    "colab": {
      "name": "CS224W - Colab 0.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/drynoise/hello-world/blob/master/CS224W_Colab_0.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XuXWJLEm2UWS"
      },
      "source": [
        "# **CS224W - Colab 0**\n",
        "\n",
        "Colab 0 **will not be graded**, so you don't need to hand in this notebook. That said, we highly recommend you to run this notebook, so you can get familiar with the basic concepts of graph mining and Graph Neural Networks.\n",
        "\n",
        "In this Colab, we will introduce two packages, [NetworkX](https://networkx.org/documentation/stable/) and [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/).\n",
        "\n",
        "For the PyTorch Geometric section, you don't need to understand all the details already. Concepts and implementations of graph neural network will be covered in future lectures and Colabs.\n",
        "\n",
        "Please make a copy before you proceed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AIUjFBAcdaHX"
      },
      "source": [
        "# New Section"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8gzsP50bF6Gb"
      },
      "source": [
        "\n",
        "# NetworkX Tutorial\n",
        "\n",
        "NetworkX is one of the most frequently used Python packages to create, manipulate, and mine graphs.\n",
        "\n",
        "Main parts of this tutorial are adapted from https://colab.research.google.com/github/jdwittenauer/ipython-notebooks/blob/master/notebooks/libraries/NetworkX.ipynb#scrollTo=zA1OO6huHeV6"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nwwq0nSdmsOL"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QeqN7MHvH1OA"
      },
      "source": [
        "# Import the NetworkX package\n",
        "import networkx as nx"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCIeGAfLfAMK"
      },
      "source": [
        "## Graph\n",
        "NetworkX provides several classes to store different types of graphs, such as directed and undirected graph. It also provides classes to create multigraphs (both directed and undirected).\n",
        "\n",
        "For more information, please refer to [NetworkX graph types](https://networkx.org/documentation/stable/reference/classes/index.html)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qSI_n6P-e0PW",
        "outputId": "d0ab4544-4b65-41f0-db9c-0852c4f4f761",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 236
        }
      },
      "source": [
        "# Create an undirected graph G\n",
        "G = nx.Graph()\n",
        "print(G.is_directed())\n",
        "\n",
        "# Create a directed graph H\n",
        "H = nx.DiGraph()\n",
        "print(H.is_directed())\n",
        "\n",
        "# Add graph level attribute\n",
        "G.graph[\"Name\"] = \"Bar\"\n",
        "print(G.graph)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-1-017b894b428e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Create an undirected graph G\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mG\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGraph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mG\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_directed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;31m# Create a directed graph H\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'nx' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x0pLs0-Ka9j8"
      },
      "source": [
        "## Node\n",
        "\n",
        "Nodes (with attributes) can be easily added to NetworkX graphs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jQ8ApAL5H1GB"
      },
      "source": [
        "# Add one node with node level attributes\n",
        "G.add_node(0, feature=0, label=0)\n",
        "\n",
        "# Get attributes of the node 0\n",
        "node_0_attr = G.nodes[0]\n",
        "print(\"Node 0 has the attributes {}\".format(node_0_attr))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "btOdMYqJaKia"
      },
      "source": [
        "# Add multiple nodes with attributes\n",
        "G.add_nodes_from([\n",
        "  (1, {\"feature\": 1, \"label\": 1}),\n",
        "  (2, {\"feature\": 2, \"label\": 2})\n",
        "])\n",
        "\n",
        "# Loop through all the nodes\n",
        "# Set data=True will return node attributes\n",
        "for node in G.nodes(data=True):\n",
        "  print(node)\n",
        "\n",
        "# Get number of nodes\n",
        "num_nodes = G.number_of_nodes()\n",
        "print(\"G has {} nodes\".format(num_nodes))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0AdoaZPgbRis"
      },
      "source": [
        "## Edge\n",
        "\n",
        "Similar to nodes, edges (with attributes) can also be easily added to NetworkX graphs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0szH5F6EH079"
      },
      "source": [
        "# Add one edge with edge weight 0.5\n",
        "G.add_edge(0, 1, weight=0.5)\n",
        "\n",
        "# Get attributes of the edge (0, 1)\n",
        "edge_0_1_attr = G.edges[(0, 1)]\n",
        "print(\"Edge (0, 1) has the attributes {}\".format(edge_0_1_attr))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QRTmi4EUaf_I"
      },
      "source": [
        "# Add multiple edges with edge weights\n",
        "G.add_edges_from([\n",
        "  (1, 2, {\"weight\": 0.3}),\n",
        "  (2, 0, {\"weight\": 0.1})\n",
        "])\n",
        "\n",
        "# Loop through all the edges\n",
        "# Here there is no data=True, so only the edge will be returned\n",
        "for edge in G.edges():\n",
        "  print(edge)\n",
        "\n",
        "# Get number of edges\n",
        "num_edges = G.number_of_edges()\n",
        "print(\"G has {} edges\".format(num_edges))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9u1Utjn4bc7k"
      },
      "source": [
        "## Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lykPzFOEahuP"
      },
      "source": [
        "# Draw the graph\n",
        "nx.draw(G, with_labels = True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Q6YTP2FDbOS"
      },
      "source": [
        "## Node Degree and Neighbor"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GFA3B6Z_DE3q"
      },
      "source": [
        "node_id = 1\n",
        "\n",
        "# Degree of node 1\n",
        "print(\"Node {} has degree {}\".format(node_id, G.degree[node_id]))\n",
        "\n",
        "# Get neighbor of node 1\n",
        "for neighbor in G.neighbors(node_id):\n",
        "  print(\"Node {} has neighbor {}\".format(node_id, neighbor))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4gVRVckZeSdA"
      },
      "source": [
        "## Other Functionalities\n",
        "\n",
        "NetworkX also provides plenty of useful methods to study graphs.\n",
        "\n",
        "Here is an example to get [PageRank](https://networkx.org/documentation/stable/reference/algorithms/generated/networkx.algorithms.link_analysis.pagerank_alg.pagerank.html#networkx.algorithms.link_analysis.pagerank_alg.pagerank) of nodes (we will talk about PageRank in one of the future lectures)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_gZfQ82Wiuvv"
      },
      "source": [
        "num_nodes = 4\n",
        "# Create a new path like graph and change it to a directed graph\n",
        "G = nx.DiGraph(nx.path_graph(num_nodes))\n",
        "nx.draw(G, with_labels = True)\n",
        "\n",
        "# Get the PageRank\n",
        "pr = nx.pagerank(G, alpha=0.8)\n",
        "pr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RrIFCJnlvGkg"
      },
      "source": [
        "## Documentation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_7PBwhIKu3et"
      },
      "source": [
        "You can explore more NetworkX functions through its [documentation](https://networkx.org/documentation/stable/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MDcCjrW3JWzG"
      },
      "source": [
        "# PyTorch Geometric Tutorial\n",
        "\n",
        "PyTorch Geometric (PyG) is an extension library for PyTorch. It provides useful primitives to develop Graph Deep Learning models, including various graph neural network layers and a large number of benchmark datasets.\n",
        "\n",
        "Don't worry if you don't understand some concepts such as `GCNConv` -- we will cover all of them in the future lectures :)\n",
        "\n",
        "This tutorial is adapted from https://colab.research.google.com/drive/1h3-vJGRVloF5zStxL5I0rSy4ZUPNsjy8?usp=sharing#scrollTo=ci-LpZWhRJoI by [Matthias Fey](https://rusty1s.github.io/#/)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WU7NCZtQ6msa"
      },
      "source": [
        "import torch\n",
        "print(\"PyTorch has version {}\".format(torch.__version__))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n7gMjMCT677n"
      },
      "source": [
        "## Setup\n",
        "\n",
        "The installation of PyG on Colab can be a little bit tricky. Execute the cell below -- in case of issues, more information can be found on the [PyG's installation page](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tRNEKu-R66Cw"
      },
      "source": [
        "# Install torch geometric\n",
        "!pip install -q torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html\n",
        "!pip install -q torch-sparse -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html\n",
        "!pip install -q torch-geometric"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HLxnaKsN8GVf"
      },
      "source": [
        "## Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qywlcjyr8USw"
      },
      "source": [
        "# Helper function for visualization.\n",
        "%matplotlib inline\n",
        "import torch\n",
        "import networkx as nx\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Visualization function for NX graph or PyTorch tensor\n",
        "def visualize(h, color, epoch=None, loss=None):\n",
        "    plt.figure(figsize=(7,7))\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "\n",
        "    if torch.is_tensor(h):\n",
        "        h = h.detach().cpu().numpy()\n",
        "        plt.scatter(h[:, 0], h[:, 1], s=140, c=color, cmap=\"Set2\")\n",
        "        if epoch is not None and loss is not None:\n",
        "            plt.xlabel(f'Epoch: {epoch}, Loss: {loss.item():.4f}', fontsize=16)\n",
        "    else:\n",
        "        nx.draw_networkx(G, pos=nx.spring_layout(G, seed=42), with_labels=False,\n",
        "                         node_color=color, cmap=\"Set2\")\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbny-iTO7NQN"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "Recently, deep learning on graphs has emerged to one of the hottest research fields in the deep learning community.\n",
        "Here, **Graph Neural Networks (GNNs)** aim to generalize classical deep learning concepts to irregular structured data (in contrast to images or texts) and to enable neural networks to reason about objects and their relations.\n",
        "\n",
        "This tutorial will introduce you to some fundamental concepts regarding deep learning on graphs via Graph Neural Networks based on the **[PyTorch Geometric (PyG) library](https://github.com/rusty1s/pytorch_geometric)**.\n",
        "PyTorch Geometric is an extension library to the popular deep learning framework [PyTorch](https://pytorch.org/), and consists of various methods and utilities to ease the implementation of Graph Neural Networks.\n",
        "\n",
        "Following [Kipf et al. (2017)](https://arxiv.org/abs/1609.02907), let's dive into the world of GNNs by looking at a simple graph-structured example, the well-known [**Zachary's karate club network**](https://en.wikipedia.org/wiki/Zachary%27s_karate_club). This graph describes a social network of 34 members of a karate club and documents links between members who interacted outside the club. Here, we are interested in detecting communities that arise from the member's interaction."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g3uPffzbyqn9"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "PyTorch Geometric provides an easy access to the dataset via the [`torch_geometric.datasets`](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets) subpackage:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YrpL9CtS7nx2"
      },
      "source": [
        "from torch_geometric.datasets import KarateClub\n",
        "\n",
        "dataset = KarateClub()\n",
        "print(f'Dataset: {dataset}:')\n",
        "print('======================')\n",
        "print(f'Number of graphs: {len(dataset)}')\n",
        "print(f'Number of features: {dataset.num_features}')\n",
        "print(f'Number of classes: {dataset.num_classes}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lCeRGa2q7sdl"
      },
      "source": [
        "After initializing the [`KarateClub`](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.KarateClub) dataset, we first can inspect some of its properties.\n",
        "For example, we can see that this dataset holds exactly **one graph**, and that each node in this dataset is assigned a **34-dimensional feature vector** (which uniquely describes the members of the karate club).\n",
        "Furthermore, the graph holds exactly **4 classes**, which represent the community each node belongs to.\n",
        "\n",
        "Let's now look at the underlying graph in more detail:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sTLapYhP7uCn"
      },
      "source": [
        "data = dataset[0]  # Get the first graph object.\n",
        "\n",
        "print(data)\n",
        "print('==============================================================')\n",
        "\n",
        "# Gather some statistics about the graph.\n",
        "print(f'Number of nodes: {data.num_nodes}')\n",
        "print(f'Number of edges: {data.num_edges}')\n",
        "print(f'Average node degree: {data.num_edges / data.num_nodes:.2f}')\n",
        "print(f'Number of training nodes: {data.train_mask.sum()}')\n",
        "print(f'Training node label rate: {int(data.train_mask.sum()) / data.num_nodes:.2f}')\n",
        "print(f'Contains isolated nodes: {data.contains_isolated_nodes()}')\n",
        "print(f'Contains self-loops: {data.contains_self_loops()}')\n",
        "print(f'Is undirected: {data.is_undirected()}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lIzbIoc-y8J4"
      },
      "source": [
        "## Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I5zhmKIH72Rf"
      },
      "source": [
        "Each graph in PyTorch Geometric is represented by a single [`Data`](https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Data) object, which holds all the information to describe its graph representation.\n",
        "We can print the data object anytime via `print(data)` to receive a short summary about its attributes and their shapes:\n",
        "```\n",
        "Data(edge_index=[2, 156], x=[34, 34], y=[34], train_mask=[34])\n",
        "```\n",
        "We can see that this `data` object holds 4 attributes:\n",
        "(1) The `edge_index` property holds the information about the **graph connectivity**, *i.e.*, a tuple of source and destination node indices for each edge.\n",
        "PyG further refers to (2) **node features** as `x` (each of the 34 nodes is assigned a 34-dim feature vector), and to (3) **node labels** as `y` (each node is assigned to exactly one class).\n",
        "(4) There also exists an additional attribute called `train_mask`, which describes for which nodes we already know their community assigments.\n",
        "In total, we are only aware of the ground-truth labels of 4 nodes (one for each community), and the task is to infer the community assignment for the remaining nodes.\n",
        "\n",
        "The `data` object also provides some **utility functions** to infer some basic properties of the underlying graph.\n",
        "For example, we can easily infer whether there exists isolated nodes in the graph (*i.e.* there exists no edge to any node), whether the graph contains self-loops (*i.e.*, $(v, v) \\in \\mathcal{E}$), or whether the graph is undirected (*i.e.*, for each edge $(v, w) \\in \\mathcal{E}$ there also exists the edge $(w, v) \\in \\mathcal{E}$)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iFA6Xi4O79r0"
      },
      "source": [
        "from IPython.display import Javascript  # Restrict height of output cell.\n",
        "display(Javascript('''google.colab.output.setIframeHeight(0, true, {maxHeight: 300})'''))\n",
        "\n",
        "edge_index = data.edge_index\n",
        "print(edge_index.t())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nLLsT0ROzffp"
      },
      "source": [
        "## Edge Index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wQJyi9OB8dh_"
      },
      "source": [
        "By printing `edge_index`, we can further understand how PyG represents graph connectivity internally.\n",
        "We can see that for each edge, `edge_index` holds a tuple of two node indices, where the first value describes the node index of the source node and the second value describes the node index of the destination node of an edge.\n",
        "\n",
        "This representation is known as the **COO format (coordinate format)** commonly used for representing sparse matrices.\n",
        "Instead of holding the adjacency information in a dense representation $\\mathbf{A} \\in \\{ 0, 1 \\}^{|\\mathcal{V}| \\times |\\mathcal{V}|}$, PyG represents graphs sparsely, which refers to only holding the coordinates/values for which entries in $\\mathbf{A}$ are non-zero.\n",
        "\n",
        "We can further visualize the graph by converting it to the `networkx` library format, which implements, in addition to graph manipulation functionalities, powerful tools for visualization:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3KfJHtlV8h3W"
      },
      "source": [
        "from torch_geometric.utils import to_networkx\n",
        "\n",
        "G = to_networkx(data, to_undirected=True)\n",
        "visualize(G, color=data.y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUdHZY2u8vn3"
      },
      "source": [
        "## Implementing Graph Neural Networks\n",
        "\n",
        "After learning about PyG's data handling, it's time to implement our first Graph Neural Network!\n",
        "\n",
        "For this, we will use one of the most simple GNN operators, the **GCN layer** ([Kipf et al. (2017)](https://arxiv.org/abs/1609.02907)).\n",
        "\n",
        "PyG implements this layer via [`GCNConv`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GCNConv), which can be executed by passing in the node feature representation `x` and the COO graph connectivity representation `edge_index`.\n",
        "\n",
        "With this, we are ready to create our first Graph Neural Network by defining our network architecture in a `torch.nn.Module` class:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6tQGQF8r8zIr"
      },
      "source": [
        "import torch\n",
        "from torch.nn import Linear\n",
        "from torch_geometric.nn import GCNConv\n",
        "\n",
        "\n",
        "class GCN(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super(GCN, self).__init__()\n",
        "        torch.manual_seed(12345)\n",
        "        self.conv1 = GCNConv(dataset.num_features, 4)\n",
        "        self.conv2 = GCNConv(4, 4)\n",
        "        self.conv3 = GCNConv(4, 2)\n",
        "        self.classifier = Linear(2, dataset.num_classes)\n",
        "\n",
        "    def forward(self, x, edge_index):\n",
        "        h = self.conv1(x, edge_index)\n",
        "        h = h.tanh()\n",
        "        h = self.conv2(h, edge_index)\n",
        "        h = h.tanh()\n",
        "        h = self.conv3(h, edge_index)\n",
        "        h = h.tanh()  # Final GNN embedding space.\n",
        "        \n",
        "        # Apply a final (linear) classifier.\n",
        "        out = self.classifier(h)\n",
        "\n",
        "        return out, h\n",
        "\n",
        "model = GCN()\n",
        "print(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5zgbaD5P8_M_"
      },
      "source": [
        "Here, we first initialize all of our building blocks in `__init__` and define the computation flow of our network in `forward`.\n",
        "We first define and stack **three graph convolution layers**, which corresponds to aggregating 3-hop neighborhood information around each node (all nodes up to 3 \"hops\" away).\n",
        "In addition, the `GCNConv` layers reduce the node feature dimensionality to $2$, *i.e.*, $34 \\rightarrow 4 \\rightarrow 4 \\rightarrow 2$. Each `GCNConv` layer is enhanced by a [tanh](https://pytorch.org/docs/stable/generated/torch.nn.Tanh.html?highlight=tanh#torch.nn.Tanh) non-linearity.\n",
        "\n",
        "After that, we apply a single linear transformation ([`torch.nn.Linear`](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html?highlight=linear#torch.nn.Linear)) that acts as a classifier to map our nodes to 1 out of the 4 classes/communities.\n",
        "\n",
        "We return both the output of the final classifier as well as the final node embeddings produced by our GNN.\n",
        "We proceed to initialize our final model via `GCN()`, and printing our model produces a summary of all its used sub-modules."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "48uhs_0j9AMX"
      },
      "source": [
        "model = GCN()\n",
        "\n",
        "_, h = model(data.x, data.edge_index)\n",
        "print(f'Embedding shape: {list(h.shape)}')\n",
        "\n",
        "visualize(h, color=data.y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nDtJ9Zjw9I_Y"
      },
      "source": [
        "Remarkably, even before training the weights of our model, the model produces an embedding of nodes that closely resembles the community-structure of the graph.\n",
        "Nodes of the same color (community) are already closely clustered together in the embedding space, although the weights of our model are initialized **completely at random** and we have not yet performed any training so far!\n",
        "This leads to the conclusion that GNNs introduce a strong inductive bias, leading to similar embeddings for nodes that are close to each other in the input graph.\n",
        "\n",
        "### Training on the Karate Club Network\n",
        "\n",
        "But can we do better? Let's look at an example on how to train our network parameters based on the knowledge of the community assignments of 4 nodes in the graph (one for each community):\n",
        "\n",
        "Since everything in our model is differentiable and parameterized, we can add some labels, train the model and observe how the embeddings react.\n",
        "Here, we make use of a semi-supervised or transductive learning procedure: We simply train against one node per class, but are allowed to make use of the complete input graph data.\n",
        "\n",
        "Training our model is very similar to any other PyTorch model.\n",
        "In addition to defining our network architecture, we define a loss critertion (here, [`CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)) and initialize a stochastic gradient optimizer (here, [`Adam`](https://pytorch.org/docs/stable/optim.html?highlight=adam#torch.optim.Adam)).\n",
        "After that, we perform multiple rounds of optimization, where each round consists of a forward and backward pass to compute the gradients of our model parameters w.r.t. to the loss derived from the forward pass.\n",
        "If you are not new to PyTorch, this scheme should appear familar to you. \n",
        "Otherwise, the PyTorch docs provide [a good introduction on how to train a neural network in PyTorch](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#define-a-loss-function-and-optimizer).\n",
        "\n",
        "Note that our semi-supervised learning scenario is achieved by the following line:\n",
        "```\n",
        "loss = criterion(out[data.train_mask], data.y[data.train_mask])\n",
        "```\n",
        "While we compute node embeddings for all of our nodes, we **only make use of the training nodes for computing the loss**.\n",
        "Here, this is implemented by filtering the output of the classifier `out` and ground-truth labels `data.y` to only contain the nodes in the `train_mask`.\n",
        "\n",
        "Let us now start training and see how our node embeddings evolve over time (best experienced by explicitely running the code):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FI3DETGi9ND6"
      },
      "source": [
        "import time\n",
        "from IPython.display import Javascript  # Restrict height of output cell.\n",
        "display(Javascript('''google.colab.output.setIframeHeight(0, true, {maxHeight: 430})'''))\n",
        "\n",
        "model = GCN()\n",
        "criterion = torch.nn.CrossEntropyLoss()  # Define loss criterion.\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # Define optimizer.\n",
        "\n",
        "def train(data):\n",
        "    optimizer.zero_grad()  # Clear gradients.\n",
        "    out, h = model(data.x, data.edge_index)  # 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, h\n",
        "\n",
        "for epoch in range(401):\n",
        "    loss, h = train(data)\n",
        "    # Visualize the node embeddings every 10 epochs\n",
        "    if epoch % 10 == 0:\n",
        "        visualize(h, color=data.y, epoch=epoch, loss=loss)\n",
        "        time.sleep(0.3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F2B3X6tf9YpS"
      },
      "source": [
        "As one can see, our 3-layer GCN model manages to linearly separating the communities and classifying most of the nodes correctly.\n",
        "\n",
        "Furthermore, we did this all with a few lines of code, thanks to the PyTorch Geometric library which helped us out with data handling and GNN implementations.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9bELRjibIRO"
      },
      "source": [
        "## Documentation\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OmqyWVNObNcK"
      },
      "source": [
        "You can explore more PyG functions through its [documentation](https://pytorch-geometric.readthedocs.io/en/latest/)."
      ]
    }
  ]
}