Heterogeneous_Node_Classification_with_DeepSNAP.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "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"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XuXWJLEm2UWS"
      },
      "source": [
        "# **Heterogeneous Node Classification with DeepSNAP**\n",
        "\n",
        "Heterogeneous graphs extend the traditional homogenous graphs by specifically incorperating different node and edge types. This additional information allows us to extend the traditional graph neural nework models, such as applying the heterogenous message passing, where different message types now exist between different node, edge type relationships.\n",
        "\n",
        "In this tutorial, we will build a heterogenous graph neural netowrk model by using [PyTorch Geonetric](https://pytorch-geometric.readthedocs.io/en/latest/) and [DeepSNAP](https://snap.stanford.edu/deepsnap/) on the heterogeneous node property prediction (node classification) task."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MSaetj53YnT6"
      },
      "source": [
        "# Device\n",
        "\n",
        "Please click `Runtime` and then `Change runtime type`. Then set the `hardware accelerator` to **GPU**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67gOQITlCNQi"
      },
      "source": [
        "# Installation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "J_m9l6OYCQZP",
        "outputId": "72f0c4da-cd7b-47f1-da39-6913d3f02df1"
      },
      "source": [
        "!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\n",
        "!pip install -q git+https://github.com/snap-stanford/deepsnap.git"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/108.0 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━\u001b[0m \u001b[32m81.9/108.0 kB\u001b[0m \u001b[31m2.4 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m108.0/108.0 kB\u001b[0m \u001b[31m2.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5LsVSRuI3hU"
      },
      "source": [
        "# Heterogeneous Graph Node Classification\n",
        "\n",
        "In this tutorial, we will use PyTorch Geometric and DeepSNAP to implement a GNN model for heterogeneous graph node property prediction (node classification).\n",
        "\n",
        "At first let's take look at the general structure of a heterogeneous layer by an example.\n",
        "\n",
        "Let's assume we have a graph $G$, which contains two node types $a$ and $b$, and three message types $m_1=(a, r_1, a)$, $m_2=(a, r_2, b)$ and $m_3=(a, r_3, b)$.\n",
        "\n",
        "Thus, for $G$ a heterogeneous layer will contains three Heterogeneous GNN layers (`HeteroGNNConv` in this Colab) where each `HeteroGNNConv` layer will perform the message passing and aggregation with respect to only one message type. The overview of the heterogeneous layer is shown below:\n",
        "\n",
        "<br/>\n",
        "<center>\n",
        "<img src=\"https://web.stanford.edu/class/cs224w/images/colab4/hetero_conv.png\"/>\n",
        "</center>\n",
        "<br/>\n",
        "\n",
        "In this Colab, all the $l^{th}$ Heterogeneous GNN layers will be managed by a ($l^{th}$) Heterogeneous GNN Wrapper layer (the `HeteroGNNWrapperConv`). The $l^{th}$ Heterogeneous GNN Wrapper layer will take in the input node embeddings from $(l-1)^{th}$ layer and aggregate (across message types) the Heterogeneous GNN layers' results. For example, the wrapper layer will aggregate node type $b$'s node embeddings from Heterogeneous GNN layers for $m_2$ and $m_3$. The \"simplified\" heterogeneous layer structure is shown below:\n",
        "\n",
        "<br/>\n",
        "<center>\n",
        "<img src=\"http://web.stanford.edu/class/cs224w/images/colab4/hetero_conv_1.png\"/>\n",
        "</center>\n",
        "<br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pkFjcktiJJLm"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NAm9_OcJJJ-W"
      },
      "source": [
        "import copy\n",
        "import torch\n",
        "import deepsnap\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch_geometric.nn as pyg_nn\n",
        "\n",
        "from sklearn.metrics import f1_score\n",
        "from deepsnap.hetero_gnn import forward_op\n",
        "from deepsnap.hetero_graph import HeteroGraph\n",
        "from torch_sparse import SparseTensor, matmul"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBlboS5kJmJL"
      },
      "source": [
        "## Heterogeneous GNN Layer\n",
        "\n",
        "Now let's start working on our own implementation of a heterogeneous layer (the `HeteroGNNConv`)! In general, our heterogeneous GNN layer draws ideas from the **GraphSAGE** ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)).\n",
        "\n",
        "At first, let's implement the GNN layer for each message type:\n",
        "\n",
        "\\begin{equation}\n",
        "m =(s, r, d)\n",
        "\\end{equation}\n",
        "\n",
        "Each message type is a tuple containing three elements where $s$ refers to the source node type, $r$ refers to the edge (relation) type and $d$ refers to the destination node type. The update rule is very similar to that of GraphSAGE but we need to include the node types and the edge type. The update rule is described as below:\n",
        "\n",
        "\\begin{equation}\n",
        "h_v^{(l)[m]} = W^{(l)[m]} \\cdot \\text{CONCAT} \\Big( W_d^{(l)[m]} \\cdot h_v^{(l-1)}, W_s^{(l)[m]} \\cdot AGG(\\{h_u^{(l-1)}, \\forall u \\in N_{m}(v) \\})\\Big)\n",
        "\\end{equation}\n",
        "\n",
        "where $[m]$ indicates that the weight matrices or embeddings with respect to message type $m$, $W_s^{(l)[m]}$ computes the messages from neighboring nodes, $W_d^{(l)[m]}$ compute messages from the node itself, and $W^{(l)[m]}$ aggregates messages from both node types. In the equation above, $v$ has the node type $d$, and $u$ has the node type $s$.\n",
        "\n",
        "For simplicity, we use mean aggregations for $AGG$ where:\n",
        "\n",
        "\\begin{equation}\n",
        "AGG(\\{h_u^{(l-1)}, \\forall u \\in N_{m}(v) \\}) = \\frac{1}{|N_{m}(v)|} \\sum_{u\\in N_{m}(v)} h_u^{(l-1)}\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_Z1b0Mf8Jova"
      },
      "source": [
        "class HeteroGNNConv(pyg_nn.MessagePassing):\n",
        "    def __init__(self, in_channels_src, in_channels_dst, out_channels):\n",
        "        super(HeteroGNNConv, self).__init__(aggr=\"mean\")\n",
        "\n",
        "        self.in_channels_src = in_channels_src\n",
        "        self.in_channels_dst = in_channels_dst\n",
        "        self.out_channels = out_channels\n",
        "\n",
        "        self.lin_src = nn.Linear(self.in_channels_src, self.out_channels)\n",
        "        self.lin_dst = nn.Linear(self.in_channels_dst, self.out_channels)\n",
        "        self.lin_update = nn.Linear(self.out_channels * 2, self.out_channels)\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        node_feature_src,\n",
        "        node_feature_dst,\n",
        "        edge_index,\n",
        "        size=None\n",
        "    ):\n",
        "        return self.propagate(\n",
        "            edge_index, size=size,\n",
        "            node_feature_dst=node_feature_dst,\n",
        "            node_feature_src=node_feature_src\n",
        "        )\n",
        "\n",
        "    def message_and_aggregate(self, edge_index, node_feature_src):\n",
        "        out = matmul(edge_index, node_feature_src, reduce=\"mean\")\n",
        "        return out\n",
        "\n",
        "    def update(self, aggr_out, node_feature_dst):\n",
        "        aggr_out = self.lin_src(aggr_out)\n",
        "        node_feature_dst = self.lin_dst(node_feature_dst)\n",
        "        aggr_out = torch.cat([aggr_out, node_feature_dst], dim=-1)\n",
        "        aggr_out = self.lin_update(aggr_out)\n",
        "        return aggr_out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKq8ScTiJthn"
      },
      "source": [
        "## Heterogeneous GNN Wrapper Layer\n",
        "\n",
        "After implementing the GNN layer for each message type, we need to aggregate the the node embedding results (with respect to each message types) together. Here we will implement two types of message type level aggregation.\n",
        "\n",
        "The first one is simply the mean aggregation:\n",
        "\n",
        "\\begin{equation}\n",
        "h_v^{(l)} = \\frac{1}{M}\\sum_{m=1}^{M}h_v^{(l)[m]}\n",
        "\\end{equation}\n",
        "\n",
        "Here node $v$ has the node type $d$ and $M$ is the total number of message types that the destination node type is $d$.\n",
        "\n",
        "The other one is the semantic level attention introduced in **HAN** ([Wang et al. (2019)](https://arxiv.org/abs/1903.07293)). Instead of directly averaging on the message type aggregation results, we can use attention to learn which message type result can be more important, then aggregate from all the message types. Following are the equations for semantic level attention:\n",
        "\n",
        "\\begin{equation}\n",
        "e_{m} = \\frac{1}{|V_{d}|} \\sum_{v \\in V_{d}} q_{attn}^T \\cdot tanh \\Big( W_{attn}^{(l)} \\cdot h_v^{(l)[m]} + b \\Big)\n",
        "\\end{equation}\n",
        "\n",
        "where $m$ refers to message type and $d$ refers to the destination node type. Then we can compute the attention and update the $h_v^{(l)}$:\n",
        "\n",
        "\\begin{equation}\n",
        "\\alpha_{m} = \\frac{\\exp(e_{m})}{\\sum_{m=1}^M \\exp(e_{m})}\n",
        "\\end{equation}\n",
        "\n",
        "\\begin{equation}\n",
        "h_v^{(l)} = \\sum_{m=1}^{M} \\alpha_{m} \\cdot h_v^{(l)[m]}\n",
        "\\end{equation}\n",
        "\n",
        "**Notice**: You can directly use `deepsnap.hetero_gnn.HeteroConv` directly for the mean aggregation. Here we overide the `HeteroConv` in order to support the semantic level attention."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0_bun02xJwFm"
      },
      "source": [
        "class HeteroGNNWrapperConv(deepsnap.hetero_gnn.HeteroConv):\n",
        "    def __init__(self, convs, args, aggr=\"mean\"):\n",
        "        super(HeteroGNNWrapperConv, self).__init__(convs, None)\n",
        "        self.aggr = aggr\n",
        "\n",
        "        # Map the index and message type\n",
        "        self.mapping = {}\n",
        "\n",
        "        # A numpy array that stores the final attention probability\n",
        "        self.alpha = None\n",
        "\n",
        "        if self.aggr == \"attn\":\n",
        "            self.attn_proj = nn.Sequential(\n",
        "                nn.Linear(args['hidden_size'], args['attn_size']),\n",
        "                nn.Tanh(),\n",
        "                nn.Linear(args['attn_size'], 1, bias=False)\n",
        "            )\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        super(HeteroConvWrapper, self).reset_parameters()\n",
        "        if self.aggr == \"attn\":\n",
        "            for layer in self.attn_proj.children():\n",
        "                layer.reset_parameters()\n",
        "\n",
        "    def forward(self, node_features, edge_indices):\n",
        "        message_type_emb = {}\n",
        "        for message_key, message_type in edge_indices.items():\n",
        "            src_type, edge_type, dst_type = message_key\n",
        "            node_feature_src = node_features[src_type]\n",
        "            node_feature_dst = node_features[dst_type]\n",
        "            edge_index = edge_indices[message_key]\n",
        "            message_type_emb[message_key] = (\n",
        "                self.convs[message_key](\n",
        "                    node_feature_src,\n",
        "                    node_feature_dst,\n",
        "                    edge_index,\n",
        "                )\n",
        "            )\n",
        "        node_emb = {dst: [] for _, _, dst in message_type_emb.keys()}\n",
        "        mapping = {}\n",
        "        for (src, edge_type, dst), item in message_type_emb.items():\n",
        "            mapping[len(node_emb[dst])] = (src, edge_type, dst)\n",
        "            node_emb[dst].append(item)\n",
        "        self.mapping = mapping\n",
        "        for node_type, embs in node_emb.items():\n",
        "            if len(embs) == 1:\n",
        "                node_emb[node_type] = embs[0]\n",
        "            else:\n",
        "                node_emb[node_type] = self.aggregate(embs)\n",
        "        return node_emb\n",
        "\n",
        "    def aggregate(self, xs):\n",
        "        if self.aggr == \"mean\":\n",
        "            x = torch.stack(xs, dim=-1)\n",
        "            return x.mean(dim=-1)\n",
        "        elif self.aggr == \"attn\":\n",
        "            N = xs[0].shape[0] # Number of nodes for that node type\n",
        "            M = len(xs) # Number of message types for that node type\n",
        "\n",
        "            x = torch.cat(xs, dim=0).view(M, N, -1) # M * N * D\n",
        "            z = self.attn_proj(x).view(M, N) # M * N * 1\n",
        "            z = z.mean(1) # M * 1\n",
        "            alpha = torch.softmax(z, dim=0) # M * 1\n",
        "\n",
        "            # Store the attention result to self.alpha as np array\n",
        "            self.alpha = alpha.view(-1).data.cpu().numpy()\n",
        "\n",
        "            alpha = alpha.view(M, 1, 1)\n",
        "            x = x * alpha\n",
        "            return x.sum(dim=0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tn_pnCOKJw-d"
      },
      "source": [
        "## Initialize Heterogeneous GNN Layers\n",
        "\n",
        "Now let's initialize the Heterogeneous GNN Layers. Different from homogeneous graph case, heterogeneous case can be a little bit complex.\n",
        "\n",
        "In general, we need to create a dictionary of `HeteroGNNConv` layers where the keys are message types.\n",
        "\n",
        "* To get all message types, we can use `deepsnap.hetero_graph.HeteroGraph.message_types`.\n",
        "* If we are initializing the first conv layers, we need to get the feature dimension of each node type. For this we can use `deepsnap.hetero_graph.HeteroGraph.num_node_features(node_type)` which will return the node feature dimension of `node_type`. In this function, we set each `HeteroGNNConv` `out_channels` to be `hidden_size`.\n",
        "* If we are not initializing the first conv layers, all node types will have the same embedding dimension `hidden_size` and we still set `HeteroGNNConv` `out_channels` to be `hidden_size` for simplicity.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hSBImHClJzf4"
      },
      "source": [
        "def generate_convs(hetero_graph, conv, hidden_size, first_layer=False):\n",
        "    convs = {}\n",
        "    for message_type in hetero_graph.message_types:\n",
        "        if first_layer is True:\n",
        "            src_type = message_type[0]\n",
        "            dst_type = message_type[2]\n",
        "            src_size = hetero_graph.num_node_features(src_type)\n",
        "            dst_size = hetero_graph.num_node_features(dst_type)\n",
        "            convs[message_type] = conv(src_size, dst_size, hidden_size)\n",
        "        else:\n",
        "            convs[message_type] = conv(hidden_size, hidden_size, hidden_size)\n",
        "    return convs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U39dX8EpJ3FG"
      },
      "source": [
        "## HeteroGNN\n",
        "\n",
        "Now we will make a simple HeteroGNN model which contains only two `HeteroGNNWrapperConv` layers.\n",
        "\n",
        "For the forward function in `HeteroGNN`, the model is going to be run as following:\n",
        "\n",
        "$\\text{self.convs1} \\rightarrow \\text{self.bns1} \\rightarrow \\text{self.relus1} \\rightarrow \\text{self.convs2} \\rightarrow \\text{self.bns2} \\rightarrow \\text{self.relus2} \\rightarrow \\text{self.post_mps}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rplknA8aJ6J5"
      },
      "source": [
        "class HeteroGNN(torch.nn.Module):\n",
        "    def __init__(self, hetero_graph, args, aggr=\"mean\"):\n",
        "        super(HeteroGNN, self).__init__()\n",
        "\n",
        "        self.aggr = aggr\n",
        "        self.hidden_size = args['hidden_size']\n",
        "\n",
        "        self.bns1 = nn.ModuleDict()\n",
        "        self.bns2 = nn.ModuleDict()\n",
        "        self.relus1 = nn.ModuleDict()\n",
        "        self.relus2 = nn.ModuleDict()\n",
        "        self.post_mps = nn.ModuleDict()\n",
        "\n",
        "        convs1 = generate_convs(hetero_graph, HeteroGNNConv, self.hidden_size, first_layer=True)\n",
        "        convs2 = generate_convs(hetero_graph, HeteroGNNConv, self.hidden_size)\n",
        "\n",
        "        self.convs1 = HeteroGNNWrapperConv(convs1, args, aggr=self.aggr)\n",
        "        self.convs2 = HeteroGNNWrapperConv(convs2, args, aggr=self.aggr)\n",
        "\n",
        "        for node_type in hetero_graph.node_types:\n",
        "            self.bns1[node_type] = torch.nn.BatchNorm1d(self.hidden_size, eps=args['eps'])\n",
        "            self.bns2[node_type] = torch.nn.BatchNorm1d(self.hidden_size, eps=args['eps'])\n",
        "            self.post_mps[node_type] = nn.Linear(self.hidden_size, hetero_graph.num_node_labels(node_type))\n",
        "            self.relus1[node_type] = nn.LeakyReLU()\n",
        "            self.relus2[node_type] = nn.LeakyReLU()\n",
        "\n",
        "    def forward(self, node_feature, edge_index):\n",
        "        x = node_feature\n",
        "        x = self.convs1(x, edge_index)\n",
        "        x = forward_op(x, self.bns1)\n",
        "        x = forward_op(x, self.relus1)\n",
        "        x = self.convs2(x, edge_index)\n",
        "        x = forward_op(x, self.bns2)\n",
        "        x = forward_op(x, self.relus2)\n",
        "        x = forward_op(x, self.post_mps)\n",
        "        return x\n",
        "\n",
        "    def loss(self, preds, y, indices):\n",
        "        loss = 0\n",
        "        loss_func = F.cross_entropy\n",
        "        for node_type in preds:\n",
        "            idx = indices[node_type]\n",
        "            loss += loss_func(preds[node_type][idx], y[node_type][idx])\n",
        "        return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9e7q_hUJ8zB"
      },
      "source": [
        "## Training and Testing\n",
        "\n",
        "Here are the functions to train and test."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CI5Hl_5TJ_YL"
      },
      "source": [
        "def train(model, optimizer, hetero_graph, train_idx):\n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    preds = model(hetero_graph.node_feature, hetero_graph.edge_index)\n",
        "\n",
        "    loss = model.loss(preds, hetero_graph.node_label, train_idx)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    return loss.item()\n",
        "\n",
        "def test(model, graph, indices, best_model=None, best_val=0):\n",
        "    model.eval()\n",
        "    accs = []\n",
        "    for index in indices:\n",
        "        preds = model(graph.node_feature, graph.edge_index)\n",
        "        num_node_types = 0\n",
        "        micro = 0\n",
        "        macro = 0\n",
        "        for node_type in preds:\n",
        "            idx = index[node_type]\n",
        "            pred = preds[node_type][idx]\n",
        "            pred = pred.max(1)[1]\n",
        "            label_np = graph.node_label[node_type][idx].cpu().numpy()\n",
        "            pred_np = pred.cpu().numpy()\n",
        "            micro = f1_score(label_np, pred_np, average='micro')\n",
        "            macro = f1_score(label_np, pred_np, average='macro')\n",
        "            num_node_types += 1\n",
        "        # Averaging f1 score might not make sense, but in our example we only\n",
        "        # have one node type\n",
        "        micro /= num_node_types\n",
        "        macro /= num_node_types\n",
        "        accs.append((micro, macro))\n",
        "    if accs[1][0] > best_val:\n",
        "        best_val = accs[1][0]\n",
        "        best_model = copy.deepcopy(model)\n",
        "    return accs, best_model, best_val"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DpNz9B5AKBUU"
      },
      "source": [
        "args = {\n",
        "    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),\n",
        "    'hidden_size': 64,\n",
        "    'epochs': 100,\n",
        "    'weight_decay': 1e-5,\n",
        "    'lr': 0.003,\n",
        "    'attn_size': 32,\n",
        "    'eps': 1.0,\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pRHbWD4hKED8"
      },
      "source": [
        "## Dataset and Preprocessing\n",
        "\n",
        "In the next, we will load the data and create a tensor backend (without a NetworkX graph) `deepsnap.hetero_graph.HeteroGraph` object.\n",
        "\n",
        "We will use the `ACM(3025)` dataset in our node property prediction task, which is proposed in **HAN** ([Wang et al. (2019)](https://arxiv.org/abs/1903.07293)) and our dataset is extracted from [DGL](https://www.dgl.ai/)'s [ACM.mat](https://data.dgl.ai/dataset/ACM.mat).\n",
        "\n",
        "Now, let's download the extracted dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rf0_sv5kN6DD",
        "outputId": "6fa62d7f-dce5-476c-c3c6-7332701bc42e"
      },
      "source": [
        "!wget https://www.dropbox.com/s/8c3102hm4ffm092/acm.pkl"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2021-07-11 04:27:07--  https://www.dropbox.com/s/8c3102hm4ffm092/acm.pkl\n",
            "Resolving www.dropbox.com (www.dropbox.com)... 162.125.2.18, 2620:100:6017:18::a27d:212\n",
            "Connecting to www.dropbox.com (www.dropbox.com)|162.125.2.18|:443... connected.\n",
            "HTTP request sent, awaiting response... 301 Moved Permanently\n",
            "Location: /s/raw/8c3102hm4ffm092/acm.pkl [following]\n",
            "--2021-07-11 04:27:07--  https://www.dropbox.com/s/raw/8c3102hm4ffm092/acm.pkl\n",
            "Reusing existing connection to www.dropbox.com:443.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com/cd/0/inline/BSHBFbSUrfmIRd_vnw0Kd-UA4grd9PJoO6x0xCxPvWs_Zx7p2tMLbw0GxAdSJij1SgHqK5WtqlUMt01yg1K_qVio0itEtbsczNzYKW9BmG-aCIgnUpQ1UNRRrRkfx1FAm1TZ0SldfiveE0Apa0tTuI1N/file# [following]\n",
            "--2021-07-11 04:27:08--  https://uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com/cd/0/inline/BSHBFbSUrfmIRd_vnw0Kd-UA4grd9PJoO6x0xCxPvWs_Zx7p2tMLbw0GxAdSJij1SgHqK5WtqlUMt01yg1K_qVio0itEtbsczNzYKW9BmG-aCIgnUpQ1UNRRrRkfx1FAm1TZ0SldfiveE0Apa0tTuI1N/file\n",
            "Resolving uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com (uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com)... 162.125.2.15, 2620:100:6017:15::a27d:20f\n",
            "Connecting to uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com (uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com)|162.125.2.15|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: /cd/0/inline2/BSEoiKBJ2hp7-nIwVRZpl9JwQU5upftnZKUj-I8d90Oz0oW17oqJEcKo4z6Cv1RLYBbL2_ofboG3nnkGpJvATeLxsvce-TpWcp_Hybwg8cMtYdGmPLIWHkNoy4opMx1thgYsQK-gmHCaVtvvuCht_9fbtXn3dXRo0VdLIkoT5GN30WDEvXBaDG-ANTKE62IzPslzgCoQoa0GsRmGsYq4axSFzUi1NvP46mcZPjjJD72Yk0JjcuyHYfrOz_B8zBiHnmuruVfalFjgippslJczt2ThKGJHTQ7yWwVWL2AKN3d_-Q7hfyX4a6aGNPD9dT-ENJf-vB3IR1b6JftQ4xWBi83eTzf6EcuIaamBxGopRdG_N_36BjByRdGvxHfo5p1_eRg/file [following]\n",
            "--2021-07-11 04:27:08--  https://uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com/cd/0/inline2/BSEoiKBJ2hp7-nIwVRZpl9JwQU5upftnZKUj-I8d90Oz0oW17oqJEcKo4z6Cv1RLYBbL2_ofboG3nnkGpJvATeLxsvce-TpWcp_Hybwg8cMtYdGmPLIWHkNoy4opMx1thgYsQK-gmHCaVtvvuCht_9fbtXn3dXRo0VdLIkoT5GN30WDEvXBaDG-ANTKE62IzPslzgCoQoa0GsRmGsYq4axSFzUi1NvP46mcZPjjJD72Yk0JjcuyHYfrOz_B8zBiHnmuruVfalFjgippslJczt2ThKGJHTQ7yWwVWL2AKN3d_-Q7hfyX4a6aGNPD9dT-ENJf-vB3IR1b6JftQ4xWBi83eTzf6EcuIaamBxGopRdG_N_36BjByRdGvxHfo5p1_eRg/file\n",
            "Reusing existing connection to uc20821dc68f677ff91d5a6ca514.dl.dropboxusercontent.com:443.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 58421685 (56M) [application/octet-stream]\n",
            "Saving to: ‘acm.pkl’\n",
            "\n",
            "acm.pkl             100%[===================>]  55.71M   173MB/s    in 0.3s    \n",
            "\n",
            "2021-07-11 04:27:09 (173 MB/s) - ‘acm.pkl’ saved [58421685/58421685]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cc-1DuedNwFj"
      },
      "source": [
        "The original ACM dataset has three node types and two edge (relation) types. For simplicity, we simplify the heterogeneous graph to one node type and two edge types (shown below). This means that in our heterogeneous graph, we have one node type (paper) and two message types *(paper, author, paper)* and *(paper, subject, paper)*.\n",
        "\n",
        "<br/>\n",
        "<center>\n",
        "<img src=\"http://web.stanford.edu/class/cs224w/images/colab4/cs224w-acm.png\"/>\n",
        "</center>\n",
        "\n",
        "Following is the code for dataset preprocessing. Here for efficiency, we only use the tensor backend for the DeepSNAP `HeteroGraph`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lJy03_IsKGh6",
        "outputId": "1b7026b4-9b8b-48e2-a1cb-3b2c35661d22"
      },
      "source": [
        "print(\"Device: {}\".format(args['device']))\n",
        "\n",
        "# Load the data\n",
        "data = torch.load(\"acm.pkl\")\n",
        "\n",
        "# Message types\n",
        "message_type_1 = (\"paper\", \"author\", \"paper\")\n",
        "message_type_2 = (\"paper\", \"subject\", \"paper\")\n",
        "\n",
        "# Dictionary of edge indices\n",
        "edge_index = {}\n",
        "edge_index[message_type_1] = data['pap']\n",
        "edge_index[message_type_2] = data['psp']\n",
        "\n",
        "# Dictionary of node features\n",
        "node_feature = {}\n",
        "node_feature[\"paper\"] = data['feature']\n",
        "\n",
        "# Dictionary of node labels\n",
        "node_label = {}\n",
        "node_label[\"paper\"] = data['label']\n",
        "\n",
        "# Load the train, validation and test indices\n",
        "train_idx = {\"paper\": data['train_idx'].to(args['device'])}\n",
        "val_idx = {\"paper\": data['val_idx'].to(args['device'])}\n",
        "test_idx = {\"paper\": data['test_idx'].to(args['device'])}\n",
        "\n",
        "# Construct a deepsnap tensor backend HeteroGraph\n",
        "hetero_graph = HeteroGraph(\n",
        "    node_feature=node_feature,\n",
        "    node_label=node_label,\n",
        "    edge_index=edge_index,\n",
        "    directed=True\n",
        ")\n",
        "\n",
        "print(f\"ACM heterogeneous graph: {hetero_graph.num_nodes()} nodes, {hetero_graph.num_edges()} edges\")\n",
        "\n",
        "# Node feature and node label to device\n",
        "for key in hetero_graph.node_feature:\n",
        "    hetero_graph.node_feature[key] = hetero_graph.node_feature[key].to(args['device'])\n",
        "for key in hetero_graph.node_label:\n",
        "    hetero_graph.node_label[key] = hetero_graph.node_label[key].to(args['device'])\n",
        "\n",
        "# Edge_index to sparse tensor and to device\n",
        "for key in hetero_graph.edge_index:\n",
        "    edge_index = hetero_graph.edge_index[key]\n",
        "    adj = SparseTensor(row=edge_index[0], col=edge_index[1], sparse_sizes=(hetero_graph.num_nodes('paper'), hetero_graph.num_nodes('paper')))\n",
        "    hetero_graph.edge_index[key] = adj.t().to(args['device'])\n",
        "print(hetero_graph.edge_index[message_type_1])\n",
        "print(hetero_graph.edge_index[message_type_2])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Device: cuda\n",
            "ACM heterogeneous graph: {'paper': 3025} nodes, {('paper', 'author', 'paper'): 26256, ('paper', 'subject', 'paper'): 2207736} edges\n",
            "SparseTensor(row=tensor([   0,    0,    0,  ..., 3024, 3024, 3024], device='cuda:0'),\n",
            "             col=tensor([   8,   20,   51,  ..., 2948, 2983, 2991], device='cuda:0'),\n",
            "             size=(3025, 3025), nnz=26256, density=0.29%)\n",
            "SparseTensor(row=tensor([   0,    0,    0,  ..., 3024, 3024, 3024], device='cuda:0'),\n",
            "             col=tensor([  75,  434,  534,  ..., 3020, 3021, 3022], device='cuda:0'),\n",
            "             size=(3025, 3025), nnz=2207736, density=24.13%)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hrmU5-QQKJv6"
      },
      "source": [
        "## Start Training!\n",
        "\n",
        "Now lets start training!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0HplV9hKMkc"
      },
      "source": [
        "## Training the Mean Aggregation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZgwfyzLbKOUw",
        "outputId": "f0f2c883-2be5-4bf6-849e-33c98cf02e73"
      },
      "source": [
        "best_model = None\n",
        "best_val = 0\n",
        "\n",
        "model = HeteroGNN(hetero_graph, args, aggr=\"mean\").to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])\n",
        "\n",
        "for epoch in range(args['epochs']):\n",
        "    loss = train(model, optimizer, hetero_graph, train_idx)\n",
        "    accs, best_model, best_val = test(model, hetero_graph, [train_idx, val_idx, test_idx], best_model, best_val)\n",
        "    print(\n",
        "        f\"Epoch {epoch + 1}: loss {round(loss, 5)}, \"\n",
        "        f\"train micro {round(accs[0][0] * 100, 2)}%, train macro {round(accs[0][1] * 100, 2)}%, \"\n",
        "        f\"valid micro {round(accs[1][0] * 100, 2)}%, valid macro {round(accs[1][1] * 100, 2)}%, \"\n",
        "        f\"test micro {round(accs[2][0] * 100, 2)}%, test macro {round(accs[2][1] * 100, 2)}%\"\n",
        "    )\n",
        "best_accs, _, _ = test(best_model, hetero_graph, [train_idx, val_idx, test_idx])\n",
        "print(\n",
        "    f\"Best model: \"\n",
        "    f\"train micro {round(best_accs[0][0] * 100, 2)}%, train macro {round(best_accs[0][1] * 100, 2)}%, \"\n",
        "    f\"valid micro {round(best_accs[1][0] * 100, 2)}%, valid macro {round(best_accs[1][1] * 100, 2)}%, \"\n",
        "    f\"test micro {round(best_accs[2][0] * 100, 2)}%, test macro {round(best_accs[2][1] * 100, 2)}%\"\n",
        ")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1: loss 1.09992, train micro 33.33%, train macro 16.67%, valid micro 33.33%, valid macro 16.67%, test micro 32.89%, test macro 16.5%\n",
            "Epoch 2: loss 1.09041, train micro 59.0%, train macro 55.65%, valid micro 51.0%, valid macro 42.31%, test micro 48.05%, test macro 39.46%\n",
            "Epoch 3: loss 1.05921, train micro 86.67%, train macro 86.43%, valid micro 83.33%, valid macro 83.13%, test micro 64.09%, test macro 61.22%\n",
            "Epoch 4: loss 0.99703, train micro 75.33%, train macro 71.71%, valid micro 71.67%, valid macro 66.99%, test micro 66.4%, test macro 58.07%\n",
            "Epoch 5: loss 0.89204, train micro 72.33%, train macro 66.85%, valid micro 70.0%, valid macro 63.48%, test micro 65.93%, test macro 56.62%\n",
            "Epoch 6: loss 0.73999, train micro 71.33%, train macro 64.96%, valid micro 69.67%, valid macro 62.51%, test micro 65.88%, test macro 56.07%\n",
            "Epoch 7: loss 0.55846, train micro 72.17%, train macro 66.35%, valid micro 70.33%, valid macro 63.38%, test micro 66.02%, test macro 56.17%\n",
            "Epoch 8: loss 0.38873, train micro 74.5%, train macro 69.99%, valid micro 71.0%, valid macro 64.85%, test micro 66.26%, test macro 56.51%\n",
            "Epoch 9: loss 0.26416, train micro 77.67%, train macro 74.67%, valid micro 73.0%, valid macro 68.23%, test micro 67.01%, test macro 57.88%\n",
            "Epoch 10: loss 0.18468, train micro 82.33%, train macro 80.76%, valid micro 79.0%, valid macro 76.99%, test micro 68.19%, test macro 60.17%\n",
            "Epoch 11: loss 0.13266, train micro 89.5%, train macro 89.06%, valid micro 86.0%, valid macro 85.45%, test micro 70.21%, test macro 63.85%\n",
            "Epoch 12: loss 0.0953, train micro 94.17%, train macro 94.06%, valid micro 90.0%, valid macro 89.82%, test micro 73.36%, test macro 68.87%\n",
            "Epoch 13: loss 0.06786, train micro 97.0%, train macro 96.98%, valid micro 93.0%, valid macro 92.91%, test micro 77.32%, test macro 74.62%\n",
            "Epoch 14: loss 0.04902, train micro 98.0%, train macro 97.99%, valid micro 94.0%, valid macro 93.9%, test micro 80.24%, test macro 78.49%\n",
            "Epoch 15: loss 0.03693, train micro 98.67%, train macro 98.67%, valid micro 96.33%, valid macro 96.3%, test micro 82.54%, test macro 81.39%\n",
            "Epoch 16: loss 0.02895, train micro 99.17%, train macro 99.17%, valid micro 97.33%, valid macro 97.33%, test micro 84.19%, test macro 83.41%\n",
            "Epoch 17: loss 0.02352, train micro 99.33%, train macro 99.33%, valid micro 97.33%, valid macro 97.33%, test micro 85.32%, test macro 84.75%\n",
            "Epoch 18: loss 0.01974, train micro 99.67%, train macro 99.67%, valid micro 97.33%, valid macro 97.33%, test micro 85.65%, test macro 85.16%\n",
            "Epoch 19: loss 0.01706, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.66%, test micro 85.98%, test macro 85.55%\n",
            "Epoch 20: loss 0.01507, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 97.99%, test micro 86.21%, test macro 85.83%\n",
            "Epoch 21: loss 0.01345, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 86.26%, test macro 85.89%\n",
            "Epoch 22: loss 0.01202, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 86.31%, test macro 85.95%\n",
            "Epoch 23: loss 0.0107, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 86.12%, test macro 85.77%\n",
            "Epoch 24: loss 0.00928, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 86.02%, test macro 85.66%\n",
            "Epoch 25: loss 0.00793, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 85.93%, test macro 85.57%\n",
            "Epoch 26: loss 0.00682, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 85.79%, test macro 85.43%\n",
            "Epoch 27: loss 0.00594, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 97.99%, test micro 85.6%, test macro 85.22%\n",
            "Epoch 28: loss 0.00525, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.6%, test macro 85.23%\n",
            "Epoch 29: loss 0.00465, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 85.46%, test macro 85.08%\n",
            "Epoch 30: loss 0.00404, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 85.18%, test macro 84.79%\n",
            "Epoch 31: loss 0.00349, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.67%, test micro 84.85%, test macro 84.46%\n",
            "Epoch 32: loss 0.00301, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.01%, test micro 84.61%, test macro 84.22%\n",
            "Epoch 33: loss 0.00259, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.01%, test micro 84.47%, test macro 84.08%\n",
            "Epoch 34: loss 0.00223, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.01%, test micro 84.38%, test macro 83.99%\n",
            "Epoch 35: loss 0.00193, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.67%, test micro 84.56%, test macro 84.22%\n",
            "Epoch 36: loss 0.00168, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.67%, test micro 84.52%, test macro 84.19%\n",
            "Epoch 37: loss 0.00148, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.67%, test micro 84.66%, test macro 84.36%\n",
            "Epoch 38: loss 0.00132, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.56%, test macro 84.27%\n",
            "Epoch 39: loss 0.00119, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 84.47%, test macro 84.17%\n",
            "Epoch 40: loss 0.00108, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 84.47%, test macro 84.18%\n",
            "Epoch 41: loss 0.001, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.47%, test macro 84.18%\n",
            "Epoch 42: loss 0.00092, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.38%, test macro 84.09%\n",
            "Epoch 43: loss 0.00086, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.19%, test macro 83.9%\n",
            "Epoch 44: loss 0.00081, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.09%, test macro 83.8%\n",
            "Epoch 45: loss 0.00077, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.09%, test macro 83.81%\n",
            "Epoch 46: loss 0.00073, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 84.0%, test macro 83.73%\n",
            "Epoch 47: loss 0.0007, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.86%, test macro 83.59%\n",
            "Epoch 48: loss 0.00067, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.86%, test macro 83.59%\n",
            "Epoch 49: loss 0.00064, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.86%, test macro 83.6%\n",
            "Epoch 50: loss 0.00062, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.76%, test macro 83.52%\n",
            "Epoch 51: loss 0.0006, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.76%, test macro 83.53%\n",
            "Epoch 52: loss 0.00058, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.72%, test macro 83.48%\n",
            "Epoch 53: loss 0.00056, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.81%, test macro 83.59%\n",
            "Epoch 54: loss 0.00054, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.72%, test macro 83.49%\n",
            "Epoch 55: loss 0.00053, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.58%, test macro 83.36%\n",
            "Epoch 56: loss 0.00051, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.48%, test macro 83.27%\n",
            "Epoch 57: loss 0.0005, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.58%, test macro 83.37%\n",
            "Epoch 58: loss 0.00049, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.58%, test macro 83.38%\n",
            "Epoch 59: loss 0.00048, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.76%, test macro 83.57%\n",
            "Epoch 60: loss 0.00047, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.35%, test micro 83.76%, test macro 83.57%\n",
            "Epoch 61: loss 0.00046, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.35%, test micro 83.81%, test macro 83.63%\n",
            "Epoch 62: loss 0.00045, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.35%, test micro 83.91%, test macro 83.73%\n",
            "Epoch 63: loss 0.00044, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.35%, test micro 83.91%, test macro 83.73%\n",
            "Epoch 64: loss 0.00043, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.91%, test macro 83.75%\n",
            "Epoch 65: loss 0.00042, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.91%, test macro 83.75%\n",
            "Epoch 66: loss 0.00041, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.91%, test macro 83.75%\n",
            "Epoch 67: loss 0.00041, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.86%, test macro 83.71%\n",
            "Epoch 68: loss 0.0004, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.91%, test macro 83.76%\n",
            "Epoch 69: loss 0.00039, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.86%, test macro 83.72%\n",
            "Epoch 70: loss 0.00039, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.91%, test macro 83.77%\n",
            "Epoch 71: loss 0.00038, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.86%, test macro 83.72%\n",
            "Epoch 72: loss 0.00038, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.68%\n",
            "Epoch 73: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.91%, test macro 83.78%\n",
            "Epoch 74: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.7%\n",
            "Epoch 75: loss 0.00036, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.71%\n",
            "Epoch 76: loss 0.00036, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.71%\n",
            "Epoch 77: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.71%\n",
            "Epoch 78: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.71%\n",
            "Epoch 79: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.81%, test macro 83.71%\n",
            "Epoch 80: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.76%, test macro 83.67%\n",
            "Epoch 81: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.76%, test macro 83.67%\n",
            "Epoch 82: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.72%, test macro 83.63%\n",
            "Epoch 83: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.72%, test macro 83.64%\n",
            "Epoch 84: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.72%, test macro 83.64%\n",
            "Epoch 85: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.62%, test macro 83.55%\n",
            "Epoch 86: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.68%, test micro 83.62%, test macro 83.55%\n",
            "Epoch 87: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.58%, test macro 83.51%\n",
            "Epoch 88: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.53%, test macro 83.46%\n",
            "Epoch 89: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.48%, test macro 83.42%\n",
            "Epoch 90: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.48%, test macro 83.42%\n",
            "Epoch 91: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 92: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 93: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 94: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 95: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 96: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 97: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 98: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.37%\n",
            "Epoch 99: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.39%, test macro 83.32%\n",
            "Epoch 100: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.39%, test macro 83.32%\n",
            "Best model: train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 86.26%, test macro 85.89%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JBiYvwcuKd0z"
      },
      "source": [
        "## Training the Attention Aggregation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e6na5zyQKfvi",
        "outputId": "3b148a3d-e2f3-4d22-a6e4-726b1c1c071f"
      },
      "source": [
        "best_model = None\n",
        "best_val = 0\n",
        "\n",
        "output_size = hetero_graph.num_node_labels('paper')\n",
        "model = HeteroGNN(hetero_graph, args, aggr=\"attn\").to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])\n",
        "\n",
        "for epoch in range(args['epochs']):\n",
        "    loss = train(model, optimizer, hetero_graph, train_idx)\n",
        "    accs, best_model, best_val = test(model, hetero_graph, [train_idx, val_idx, test_idx], best_model, best_val)\n",
        "    print(\n",
        "        f\"Epoch {epoch + 1}: loss {round(loss, 5)}, \"\n",
        "        f\"train micro {round(accs[0][0] * 100, 2)}%, train macro {round(accs[0][1] * 100, 2)}%, \"\n",
        "        f\"valid micro {round(accs[1][0] * 100, 2)}%, valid macro {round(accs[1][1] * 100, 2)}%, \"\n",
        "        f\"test micro {round(accs[2][0] * 100, 2)}%, test macro {round(accs[2][1] * 100, 2)}%\"\n",
        "    )\n",
        "best_accs, _, _ = test(best_model, hetero_graph, [train_idx, val_idx, test_idx])\n",
        "print(\n",
        "    f\"Best model: \"\n",
        "    f\"train micro {round(best_accs[0][0] * 100, 2)}%, train macro {round(best_accs[0][1] * 100, 2)}%, \"\n",
        "    f\"valid micro {round(best_accs[1][0] * 100, 2)}%, valid macro {round(best_accs[1][1] * 100, 2)}%, \"\n",
        "    f\"test micro {round(best_accs[2][0] * 100, 2)}%, test macro {round(best_accs[2][1] * 100, 2)}%\"\n",
        ")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1: loss 1.10215, train micro 33.33%, train macro 16.67%, valid micro 33.33%, valid macro 16.67%, test micro 35.81%, test macro 17.58%\n",
            "Epoch 2: loss 1.0936, train micro 48.0%, train macro 39.1%, valid micro 35.0%, valid macro 20.05%, test micro 38.07%, test macro 22.36%\n",
            "Epoch 3: loss 1.06288, train micro 66.33%, train macro 55.13%, valid micro 66.0%, valid macro 55.07%, test micro 63.15%, test macro 52.73%\n",
            "Epoch 4: loss 1.00176, train micro 66.33%, train macro 54.6%, valid micro 66.0%, valid macro 54.67%, test micro 65.13%, test macro 53.77%\n",
            "Epoch 5: loss 0.90104, train micro 66.0%, train macro 53.86%, valid micro 65.67%, valid macro 53.9%, test micro 64.85%, test macro 53.12%\n",
            "Epoch 6: loss 0.76788, train micro 66.0%, train macro 53.74%, valid micro 65.67%, valid macro 53.51%, test micro 64.61%, test macro 52.83%\n",
            "Epoch 7: loss 0.63097, train micro 66.17%, train macro 54.14%, valid micro 66.0%, valid macro 53.93%, test micro 65.04%, test macro 53.27%\n",
            "Epoch 8: loss 0.50856, train micro 68.33%, train macro 58.7%, valid micro 66.67%, valid macro 55.6%, test micro 65.27%, test macro 53.73%\n",
            "Epoch 9: loss 0.40103, train micro 70.33%, train macro 62.62%, valid micro 67.67%, valid macro 57.32%, test micro 65.55%, test macro 54.23%\n",
            "Epoch 10: loss 0.31062, train micro 72.33%, train macro 66.01%, valid micro 69.33%, valid macro 60.81%, test micro 66.12%, test macro 55.36%\n",
            "Epoch 11: loss 0.23957, train micro 76.0%, train macro 71.96%, valid micro 71.33%, valid macro 64.44%, test micro 66.49%, test macro 56.05%\n",
            "Epoch 12: loss 0.18576, train micro 80.67%, train macro 78.53%, valid micro 76.0%, valid macro 72.08%, test micro 67.11%, test macro 57.37%\n",
            "Epoch 13: loss 0.14433, train micro 90.17%, train macro 89.79%, valid micro 83.33%, valid macro 81.91%, test micro 68.8%, test macro 60.85%\n",
            "Epoch 14: loss 0.11135, train micro 94.5%, train macro 94.38%, valid micro 88.33%, valid macro 87.79%, test micro 71.48%, test macro 65.72%\n",
            "Epoch 15: loss 0.08482, train micro 97.0%, train macro 96.97%, valid micro 90.0%, valid macro 89.62%, test micro 73.51%, test macro 68.99%\n",
            "Epoch 16: loss 0.06455, train micro 97.83%, train macro 97.82%, valid micro 91.0%, valid macro 90.73%, test micro 75.39%, test macro 71.79%\n",
            "Epoch 17: loss 0.0501, train micro 98.5%, train macro 98.49%, valid micro 91.67%, valid macro 91.46%, test micro 77.65%, test macro 74.98%\n",
            "Epoch 18: loss 0.03983, train micro 99.17%, train macro 99.16%, valid micro 93.33%, valid macro 93.22%, test micro 79.25%, test macro 77.16%\n",
            "Epoch 19: loss 0.03214, train micro 99.33%, train macro 99.33%, valid micro 95.67%, valid macro 95.63%, test micro 81.08%, test macro 79.55%\n",
            "Epoch 20: loss 0.02611, train micro 99.5%, train macro 99.5%, valid micro 96.0%, valid macro 95.97%, test micro 82.82%, test macro 81.7%\n",
            "Epoch 21: loss 0.02114, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.31%, test micro 83.91%, test macro 83.03%\n",
            "Epoch 22: loss 0.01709, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 84.66%, test macro 83.97%\n",
            "Epoch 23: loss 0.01388, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 85.08%, test macro 84.53%\n",
            "Epoch 24: loss 0.01128, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.6%, test macro 85.13%\n",
            "Epoch 25: loss 0.00898, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.79%, test macro 85.36%\n",
            "Epoch 26: loss 0.00716, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.79%, test macro 85.38%\n",
            "Epoch 27: loss 0.00575, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.79%, test macro 85.41%\n",
            "Epoch 28: loss 0.00466, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 85.79%, test macro 85.39%\n",
            "Epoch 29: loss 0.00382, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.6%, test macro 85.21%\n",
            "Epoch 30: loss 0.00317, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.66%, test micro 85.69%, test macro 85.31%\n",
            "Epoch 31: loss 0.00267, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.66%, test micro 85.69%, test macro 85.3%\n",
            "Epoch 32: loss 0.00228, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.69%, test macro 85.32%\n",
            "Epoch 33: loss 0.00197, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.69%, test macro 85.31%\n",
            "Epoch 34: loss 0.00172, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.65%, test macro 85.27%\n",
            "Epoch 35: loss 0.00152, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.55%, test macro 85.18%\n",
            "Epoch 36: loss 0.00135, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.51%, test macro 85.14%\n",
            "Epoch 37: loss 0.00122, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.51%, test macro 85.14%\n",
            "Epoch 38: loss 0.0011, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.36%, test macro 85.01%\n",
            "Epoch 39: loss 0.00101, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.32%, test macro 84.96%\n",
            "Epoch 40: loss 0.00093, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.36%, test macro 85.02%\n",
            "Epoch 41: loss 0.00086, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 85.22%, test macro 84.88%\n",
            "Epoch 42: loss 0.0008, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.94%, test macro 84.6%\n",
            "Epoch 43: loss 0.00075, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.89%, test macro 84.58%\n",
            "Epoch 44: loss 0.0007, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.85%, test macro 84.54%\n",
            "Epoch 45: loss 0.00066, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.75%, test macro 84.46%\n",
            "Epoch 46: loss 0.00063, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.56%, test macro 84.28%\n",
            "Epoch 47: loss 0.00059, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.47%, test macro 84.19%\n",
            "Epoch 48: loss 0.00057, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.28%, test macro 84.03%\n",
            "Epoch 49: loss 0.00054, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.24%, test macro 83.98%\n",
            "Epoch 50: loss 0.00052, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.09%, test macro 83.85%\n",
            "Epoch 51: loss 0.0005, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.24%, test macro 84.0%\n",
            "Epoch 52: loss 0.00048, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 84.09%, test macro 83.86%\n",
            "Epoch 53: loss 0.00046, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 84.09%, test macro 83.87%\n",
            "Epoch 54: loss 0.00045, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 84.09%, test macro 83.88%\n",
            "Epoch 55: loss 0.00043, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 84.09%, test macro 83.88%\n",
            "Epoch 56: loss 0.00042, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 83.95%, test macro 83.73%\n",
            "Epoch 57: loss 0.00041, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 84.0%, test macro 83.79%\n",
            "Epoch 58: loss 0.0004, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 83.95%, test macro 83.74%\n",
            "Epoch 59: loss 0.00039, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 83.91%, test macro 83.7%\n",
            "Epoch 60: loss 0.00038, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 83.81%, test macro 83.62%\n",
            "Epoch 61: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 83.72%, test macro 83.52%\n",
            "Epoch 62: loss 0.00036, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.62%, test macro 83.43%\n",
            "Epoch 63: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.53%, test macro 83.34%\n",
            "Epoch 64: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.48%, test macro 83.29%\n",
            "Epoch 65: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.2%, test macro 83.03%\n",
            "Epoch 66: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.11%, test macro 82.93%\n",
            "Epoch 67: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.11%, test macro 82.94%\n",
            "Epoch 68: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.11%, test macro 82.94%\n",
            "Epoch 69: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.06%, test macro 82.91%\n",
            "Epoch 70: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 83.01%, test macro 82.86%\n",
            "Epoch 71: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.96%, test macro 82.81%\n",
            "Epoch 72: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.87%, test macro 82.73%\n",
            "Epoch 73: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.87%, test macro 82.73%\n",
            "Epoch 74: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.78%\n",
            "Epoch 75: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.73%\n",
            "Epoch 76: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.78%, test macro 82.63%\n",
            "Epoch 77: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.78%, test macro 82.63%\n",
            "Epoch 78: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.73%\n",
            "Epoch 79: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.78%\n",
            "Epoch 80: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.74%\n",
            "Epoch 81: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.79%\n",
            "Epoch 82: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.79%\n",
            "Epoch 83: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.79%\n",
            "Epoch 84: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.74%\n",
            "Epoch 85: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.74%\n",
            "Epoch 86: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.74%\n",
            "Epoch 87: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.82%, test macro 82.7%\n",
            "Epoch 88: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.96%, test macro 82.84%\n",
            "Epoch 89: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.8%\n",
            "Epoch 90: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.8%\n",
            "Epoch 91: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.8%\n",
            "Epoch 92: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.92%, test macro 82.8%\n",
            "Epoch 93: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.82%, test macro 82.7%\n",
            "Epoch 94: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.78%, test macro 82.65%\n",
            "Epoch 95: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.78%, test macro 82.65%\n",
            "Epoch 96: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.0%, test micro 82.87%, test macro 82.75%\n",
            "Epoch 97: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.82%, test macro 82.7%\n",
            "Epoch 98: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.78%, test macro 82.64%\n",
            "Epoch 99: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.78%, test macro 82.64%\n",
            "Epoch 100: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.67%, test micro 82.73%, test macro 82.6%\n",
            "Best model: train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.66%, test micro 85.69%, test macro 85.31%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQgx5y4UqMHH"
      },
      "source": [
        "## Attention for each Message Type\n",
        "\n",
        "Through message type level attention we can learn that which message type is more important to which layer.\n",
        "\n",
        "Here we will print out and show that each layer pay how much attention on each message type."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hvK58gijqN_C",
        "outputId": "cf06ed81-fd52-4976-b445-343f87e25e31"
      },
      "source": [
        "if model.convs1.alpha is not None and model.convs2.alpha is not None:\n",
        "    for idx, message_type in model.convs1.mapping.items():\n",
        "        print(f\"Layer 1 has attention {model.convs1.alpha[idx]} on message type {message_type}\")\n",
        "    for idx, message_type in model.convs2.mapping.items():\n",
        "        print(f\"Layer 2 has attention {model.convs2.alpha[idx]} on message type {message_type}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Layer 1 has attention 0.8046650886535645 on message type ('paper', 'author', 'paper')\n",
            "Layer 1 has attention 0.19533488154411316 on message type ('paper', 'subject', 'paper')\n",
            "Layer 2 has attention 0.435543030500412 on message type ('paper', 'author', 'paper')\n",
            "Layer 2 has attention 0.5644569993019104 on message type ('paper', 'subject', 'paper')\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5AGjsRqzAVri"
      },
      "source": [
        "More heterogeneous node classification examples please see the [examples/node_classification_hetero](https://github.com/snap-stanford/deepsnap/tree/master/examples/node_classification_hetero)."
      ]
    }
  ]
}