Colab4.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XuXWJLEm2UWS"
      },
      "source": [
        "# CS224W - Colab 4\n",
        "\n",
        "本次实验中，我们要将注意力从同构图转向异构图。异构图和我们之前见过的传统同构图不一样，它有多种结点类型和边类型。我们可以用这些额外信息来拓展我们之前的图神经网络模型，即基于异构图的消息传递。其中，不同类型的结点和边类型之间存在着不同的消息类型。\n",
        "\n",
        "首先，我们要学习如何将 `NetworkX` 图转化为用 `DeepSNAP` 来表示。除此之外，我们还要深入理解 `DeepSNAP` 如何将异构图存储和表示成 `PyTorch` 张量的。\n",
        "\n",
        "然后，我们要使用 `PyG` 和 `DeepSNAP` 来搭建一个用于结点分类任务的异构图神经网络。我们还使用异构的 ACM 数据集来评估我们的模型。\n",
        "\n",
        "**注意**：确保顺序执行所有代码单元，否则变量或库无法在之后的单元中使用。实验愉快 ：）\n",
        "\n",
        "\n",
        "> 在实验开始之前，请确保正确安装了 [PyG](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) 以及 [DeepSNAP](https://snap.stanford.edu/deepsnap/notes/installation.html) 第三方库。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PRfgbfTjCRD_"
      },
      "outputs": [],
      "source": [
        "import torch_geometric\n",
        "torch_geometric.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZoXlf4MtYrbz"
      },
      "source": [
        "## 1 DeepSNAP 异构图\n",
        "\n",
        "首先，我们来学习以下如何将 `NetworkX` 图转化为 `DeepSNAP` 支持的格式。\n",
        "\n",
        "为了将传统的 `DeepSNAP` 图表征技术拓展到可以表征异构图，我们加上以下几个图特征：\n",
        "\n",
        "- `node_feature`：结点的特征 `torch.tensor`\n",
        "- `edge_feature`：边的特征 `torch.tensor`\n",
        "- `node_label`：结点的标签 `int`\n",
        "- `node_type`：结点类型 `string`\n",
        "- `edge_type`：边类型 `string`\n",
        "\n",
        "在新加入的这几个特征之中，最关键的是 `node_type` 和 `edge_type`，它们将指导如何进行消息传递。\n",
        "\n",
        "在第一个问题中，我们将使用熟悉的 [karate club graph](https://networkx.github.io/documentation/stable/auto_examples/graph/plot_karate_club.html)。由于结点属于 `Mr.Hi` 和 `Officer` 这两种俱乐部其中的一种，我们便将结点的 `club` 属性当成 `node_type`。以下代码展示了如何区分 `NetworkX` 图中的结点。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8LQ_z5gcBVA1"
      },
      "outputs": [],
      "source": [
        "from pylab import *\n",
        "import networkx as nx\n",
        "from networkx.algorithms.community import greedy_modularity_communities\n",
        "import matplotlib.pyplot as plt\n",
        "import copy\n",
        "\n",
        "G = nx.karate_club_graph()\n",
        "community_map = {}\n",
        "for node, club in G.nodes(data=\"club\"):\n",
        "  if club == \"Mr. Hi\":\n",
        "    community_map[node] = 0\n",
        "  else:\n",
        "    community_map[node] = 1\n",
        "node_color = []\n",
        "color_map = {0: 0, 1: 1}\n",
        "node_color = [color_map[community_map[node]] for node in G.nodes()]\n",
        "pos = nx.spring_layout(G)\n",
        "plt.figure(figsize=(7, 7))\n",
        "nx.draw(G, pos=pos, cmap=plt.get_cmap('coolwarm'), node_color=node_color)\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1JpFb9fTw1lg"
      },
      "source": [
        "### 问题 1.1: 分配结点类型和结点特征\n",
        "\n",
        "使用 `community_map` 字典和上边的图 `G`，在图 `G` 中加入结点类型和结点标签。即给 `\"Mr. Hi\"` 俱乐部的结点的 `node_type` 赋值为 `n0`，给 `\"Officer\"` 俱乐部的结点的 `node_type` 赋值为 `n1`。\n",
        "\n",
        "至于 `node_label`，`\"Mr. Hi\"` 中结点的 `node_label` 赋值为 `0`，`\"Officer\"` 中结点赋值为 `1`。\n",
        "\n",
        "最后，给每个结点的特征都赋值为 `[1, 1, 1, 1, 1]`。\n",
        "\n",
        "**提示**：看一看 `NetworkX` 的函数 `nx.classes.function.set_node_attributes`。\n",
        "\n",
        "**注意**：这个问题对于后续问题很重要，请确保正确实现。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zev_hMJHJXK1"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "def assign_node_types(G, community_map):\n",
        "  # TODO: Implement this function that takes in a NetworkX graph\n",
        "  # G and community map assignment (mapping node ids --> 0/1 labels)\n",
        "  # and adds 'node_type' as a node_attribute in G.\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~2 line of code)\n",
        "  ## Note\n",
        "  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`\n",
        "  \n",
        "\n",
        "  #########################################\n",
        "  return\n",
        "\n",
        "def assign_node_labels(G, community_map):\n",
        "  # TODO: Implement this function that takes in a NetworkX graph\n",
        "  # G and community map assignment (mapping node ids --> 0/1 labels)\n",
        "  # and adds 'node_label' as a node_attribute in G.\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~2 line of code)\n",
        "  ## Note\n",
        "  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`\n",
        "\n",
        "  #########################################\n",
        "  return\n",
        "\n",
        "def assign_node_features(G):\n",
        "  # TODO: Implement this function that takes in a NetworkX graph\n",
        "  # G and adds 'node_feature' as a node_attribute in G. Each node\n",
        "  # in the graph has the same feature vector [1., 1., 1., 1., 1.]\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~2 line of code)\n",
        "  ## Note\n",
        "  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`\n",
        "\n",
        "  #########################################\n",
        "  return\n",
        "\n",
        "assign_node_types(G, community_map)\n",
        "assign_node_labels(G, community_map)\n",
        "assign_node_features(G)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mafN0P3EOhSb"
      },
      "source": [
        "### 问题 1.2: 分配边类型\n",
        "\n",
        "下一步，我们给边分配三种 `edge_type`：\n",
        "\n",
        "- `Mr. Hi` 内部的边类型：`e0`\n",
        "- `Officer` 内部的边类型：`e1`\n",
        "- 两个俱乐部之间的边类型：`e2`\n",
        "\n",
        "**提示**：使用之前的 `community_map` 和 `nx.classes.function.set_edge_attrbutes`。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rsbYWEVwSV5n"
      },
      "outputs": [],
      "source": [
        "def assign_edge_types(G, community_map):\n",
        "  # TODO: Implement this function that takes in a NetworkX graph\n",
        "  # G and community map assignment (mapping node ids --> 0/1 labels)\n",
        "  # and adds 'edge_type' as a edge_attribute in G.\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~5 line of code)\n",
        "  ## Note\n",
        "  ## 1. Create an edge assignment dict following rules above\n",
        "\n",
        "  #########################################\n",
        "  return\n",
        "\n",
        "assign_edge_types(G, community_map)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBsTdTPVTQ52"
      },
      "source": [
        "### 异构图可视化\n",
        "\n",
        "现在我们来可视化刚生成的异构图。其中有俱乐部内部边（2种）以及俱乐部之间的边（1种）。不同类型的结点和边用不同颜色来区分。下面的 `NetworkX` 对象 `G` 可以直接转换成 `deepsnap.hetero_graph.HeteroGraph` 对象。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n2sdufbODHtp"
      },
      "outputs": [],
      "source": [
        "edge_color = {}\n",
        "for edge in G.edges():\n",
        "  n1, n2 = edge\n",
        "  if community_map[n1] == community_map[n2] and community_map[n1] == 0:\n",
        "    edge_color[edge] = 'blue'\n",
        "  elif community_map[n1] == community_map[n2] and community_map[n1] == 1:\n",
        "    edge_color[edge] = 'red'\n",
        "  else:\n",
        "    edge_color[edge] = 'green'\n",
        "\n",
        "G_orig = copy.deepcopy(G)\n",
        "nx.classes.function.set_edge_attributes(G, edge_color, name='color')\n",
        "colors = nx.get_edge_attributes(G,'color').values()\n",
        "labels = nx.get_node_attributes(G, 'node_type')\n",
        "plt.figure(figsize=(8, 8))\n",
        "nx.draw(G, pos=pos, cmap=plt.get_cmap('coolwarm'), node_color=node_color, edge_color=colors, labels=labels, font_color='white')\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4DW8L0hxbxw4"
      },
      "source": [
        "### 转换为 `DeepSNAP` 表征\n",
        "\n",
        "下列代码种的 `NetworkX` 对象 `G` 可以直接转换为 `deepsnap.hetero_graph.HeteroGraph` 对象！"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-ZfHZ_eoVVGd"
      },
      "outputs": [],
      "source": [
        "from deepsnap.hetero_graph import HeteroGraph\n",
        "\n",
        "hete = HeteroGraph(G_orig)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "izq4t_O9WxDH"
      },
      "source": [
        "### 问题 1.3: 不同结点类型的结点各有多少？"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wRNSP6nnW78C"
      },
      "outputs": [],
      "source": [
        "def get_nodes_per_type(hete):\n",
        "  # TODO: Implement this function that takes a DeepSNAP dataset object\n",
        "  # and return the number of nodes per `node_type`.\n",
        "\n",
        "  num_nodes_n0 = 0\n",
        "  num_nodes_n1 = 0\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~2 line of code)\n",
        "  ## Note\n",
        "  ## 1. Colab autocomplete functionality might be useful.\n",
        "\n",
        "\n",
        "  #########################################\n",
        "\n",
        "  return num_nodes_n0, num_nodes_n1\n",
        "\n",
        "num_nodes_n0, num_nodes_n1 = get_nodes_per_type(hete)\n",
        "print(f\"Node type n0 has {num_nodes_n0} nodes\")\n",
        "print(f\"Node type n1 has {num_nodes_n1} nodes\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vEsHJp2ZYaE2"
      },
      "source": [
        "### 问题 1.4: 不同消息类型的消息各有多少？\n",
        "\n",
        "我们现在讨论的是异构的消息类型（即不同 `node_type` 和 `edge_type` 组合的消息类型）。比如，用 `e0` 类型连接的两个 `\"Mr.Hi\"` 结点是一种消息类型，即 (`n0`, `e0`, `n0`) 类型。在本问题中，我们要分析每种类型各有多少边。\n",
        "\n",
        "**提示**：如果你想要了解更多不同消息类型，试试调用 `hete.message_types`。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qobKuqbAYvJ7"
      },
      "outputs": [],
      "source": [
        "def get_num_message_edges(hete):\n",
        "  # TODO: Implement this function that takes a DeepSNAP dataset object\n",
        "  # and return the number of edges for each message type. \n",
        "  # You should return a list of tuples as \n",
        "  # (message_type, num_edge)\n",
        "\n",
        "  message_type_edges = []\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~2 line of code)\n",
        "  ## Note\n",
        "  ## 1. Colab autocomplete functionality might be useful.\n",
        "\n",
        "\n",
        "  #########################################\n",
        "\n",
        "  return message_type_edges\n",
        "\n",
        "message_type_edges = get_num_message_edges(hete)\n",
        "for (message_type, num_edges) in message_type_edges:\n",
        "  print(f\"Message type {message_type} has {num_edges} edges\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bjMVik1JbJ76"
      },
      "source": [
        "### 问题 1.5: 不同数据集分割中各有多少结点？\n",
        "\n",
        "`DeepSNAP` 中有异构图的数据集创建和分割方法。我们在这里创建用来结点分类的训练集、验证集，以及测试集，并观察一下分割后的子图。具体点就是写一个函数来计算不同数据集分割中的结点个数。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ct10Oh4gcqgD"
      },
      "outputs": [],
      "source": [
        "from deepsnap.dataset import GraphDataset\n",
        "\n",
        "def compute_dataset_split_counts(datasets):\n",
        "  # TODO: Implement this function that takes a dict of datasets\n",
        "  # and returns a dict mapping dataset names to the number of labeled\n",
        "  # nodes used for supervision in that respective dataset.  \n",
        "  \n",
        "  data_set_splits = {}\n",
        "\n",
        "  ############# Your code here ############\n",
        "  ## (~3 line of code)\n",
        "  ## Note\n",
        "  ## 1. DeepSNAP `node_label_index` will be helpful.\n",
        "  ## 2. Remember to count both node_types\n",
        "\n",
        "\n",
        "  #########################################\n",
        "\n",
        "  return data_set_splits\n",
        "\n",
        "\n",
        "dataset = GraphDataset([hete], task='node')\n",
        "# Splitting the dataset\n",
        "dataset_train, dataset_val, dataset_test = dataset.split(transductive=True, split_ratio=[0.4, 0.3, 0.3])\n",
        "datasets = {'train': dataset_train, 'val': dataset_val, 'test': dataset_test}\n",
        "\n",
        "data_set_splits = compute_dataset_split_counts(datasets)\n",
        "for dataset_name, num_nodes in data_set_splits.items():\n",
        "  print(f\"{dataset_name} dataset has {num_nodes} nodes\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lFY2PaDbVKe4"
      },
      "source": [
        "### `DeepSNAP` 数据集可视化"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iiyEw-agbgV8"
      },
      "outputs": [],
      "source": [
        "from deepsnap.dataset import GraphDataset\n",
        "\n",
        "dataset = GraphDataset([hete], task='node')\n",
        "# Splitting the dataset\n",
        "dataset_train, dataset_val, dataset_test = dataset.split(transductive=True, split_ratio=[0.4, 0.3, 0.3])\n",
        "titles = ['Train', 'Validation', 'Test']\n",
        "\n",
        "for i, dataset in enumerate([dataset_train, dataset_val, dataset_test]):\n",
        "  n0 = hete._convert_to_graph_index(dataset[0].node_label_index['n0'], 'n0').tolist()\n",
        "  n1 = hete._convert_to_graph_index(dataset[0].node_label_index['n1'], 'n1').tolist()\n",
        "\n",
        "  plt.figure(figsize=(7, 7))\n",
        "  plt.title(titles[i])\n",
        "  nx.draw(G_orig, pos=pos, node_color=\"grey\", edge_color=colors, labels=labels, font_color='white')\n",
        "  nx.draw_networkx_nodes(G_orig.subgraph(n0), pos=pos, node_color=\"blue\")\n",
        "  nx.draw_networkx_nodes(G_orig.subgraph(n1), pos=pos, node_color=\"red\")\n",
        "  show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5LsVSRuI3hU"
      },
      "source": [
        "## 2 异构图结点分类\n",
        "\n",
        "这一部分中我们将使用 `PyG` 和 `DeepSNAP` 来实现一个用来给异构图进行结点分类的 GNN 模型。本部分内容要求你对异构图有较好的理解，并熟悉如何使用 `PyG` 来实现一个 GNN 层。\n",
        "\n",
        "首先我们来看看异构层的一般结构。假设我们有一张图 $G$，其中由两种结点 $a$ 和 $b$，三种消息类型 $m_1=(a, r_1, a)$，$m_2=(a, r_2, b)$ 以及 $m_3=(a, r_3, b)$。\n",
        "\n",
        "因此，对于图 $G$ 而言，一个异构层会有三个异构 GNN 层（即本实验中的 `HeteroGNNConv`），其中每个 `HeteroGNNConv` 层都会针对一种消息类型进行消息传递与聚合。异构层概览如下：\n",
        "\n",
        "<br/>\n",
        "<center>\n",
        "<img src=\"assets/hetero_conv.png\"/>\n",
        "</center>\n",
        "<br/>\n",
        "\n",
        "本实验中，所有 $l^\\text{th}$ 的异构 GNN 层（`HeteroGNNConv`）都由 $l^\\text{th}$ 的异构 GNN 包装器层（`HeteroGNNWrapperConv`）管理。$l^\\text{th}$ 异构 GNN 包装器层会接收前一层输出的结点嵌入作为输入，并用不同的异构 GNN 层在不同消息类型间进行消息传递与聚合。比如，结点 $b$ 的嵌入就是包装器层聚合两个异构 GNN 层的结果，一个用来传播 $m_2$，一个用来传播 $m_3$。简化过的异构层结构如下：\n",
        "\n",
        "<br/>\n",
        "<center>\n",
        "<img src=\"assets/hetero_conv_1.png\"/>\n",
        "</center>\n",
        "<br/>\n",
        "\n",
        "<font color='red'>推荐你按如下步骤实现异构GNN模型：</font>\n",
        "\n",
        "1. 实现 `HeteroGNNConv` 层。\n",
        "2. 实现 `HeteroGNNWrapperConv` 中的平均聚合（`mean`聚合）。\n",
        "3. 实现 `generate_convs` 函数。\n",
        "4. 实现 `HeteroGNN` 模型和 `train` 函数。\n",
        "5. 用 `mean` 聚合训练模型聚合不同类型的消息，并确保模型性能合理。\n",
        "6. 实现 `HeteroGNNWrapperConv` 中的注意力机制聚合（`attn` 聚合）。\n",
        "7. 用 `attn` 聚合训练模型聚合不同类型的消息，并确保模型性能合理。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pkFjcktiJJLm"
      },
      "source": [
        "### 配置"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NAm9_OcJJJ-W"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBlboS5kJmJL"
      },
      "source": [
        "### 异构 GNN 层\n",
        "\n",
        "现在我们开始实现异构层（`HeteroGNNConv`）！与 Colab3 相似，我们还是使用 `PyG` 来搭建网络。总的来说，我们的异构 GNN 层从 **GraphSAGE** 中汲取了灵感([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216))。\n",
        "\n",
        "首先，我们给如下每一种消息类型都实现一个 GNN 层：\n",
        "\n",
        "$$\n",
        "\\begin{equation*}\n",
        "m =(s, r, d)\n",
        "\\end{equation*}\n",
        "$$\n",
        "\n",
        "每种消息类型都是一个三元组，其中 $s$ 指源结点类型，$r$ 指边（关系）类型，$d$ 指目标结点类型。更新规则与 GraphSAGE 相似，但我们需要将不同的结点类型和边类型考虑在内，更新规则如下：\n",
        "\n",
        "$$\n",
        "\\begin{equation*}\n",
        "h_v^{(l)[m]} = W^{(l)[m]} \\cdot \\text{CONCAT} \\left( W_d^{(l)[m]} \\cdot h_v^{(l-1)}, W_s^{(l)[m]} \\cdot \\text{AGG}\\left(\\left\\{h_u^{(l-1)}, \\forall u \\in N_{m}(v) \\right\\}\\right)\\right)\n",
        "\\end{equation*}\n",
        "$$\n",
        "\n",
        "其中 $[m]$ 指关于消息类型 $m$ 的权重矩阵或嵌入，$W_s^{(l)[m]}$ 用来计算来自邻居结点的消息，$W_d^{(l)[m]}$ 计算自己的消息，$W^{(l)[m]}$ 将两种结点类型的消息聚合。\n",
        "\n",
        "在上述等式中，$v$ 的结点类型是 $d$，$u$ 的结点类型是 $s$。简单起见，我们这里先用平均聚合作为 $\\text{AGG}$：\n",
        "\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*}\n",
        "$$\n",
        "\n",
        "> 译者注：最终，如果写成矩阵乘法的形式（聚合方式为平均聚合），则公式如下：\n",
        ">\n",
        "> $$\n",
        "> \\text{CONCAT} \\left( H^{(l-1)} W^{(l)[m]}_\\text{d}, D^{-1} A H^{(l-1)} W^{(l)[m]}_\\text{s} \\right) W^{(l)[m]}\n",
        "> $$\n",
        "> \n",
        "> 其中，$D$ 是度矩阵，$A$ 为 $m$ 类型关系的邻接矩阵。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Z1b0Mf8Jova"
      },
      "outputs": [],
      "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",
        "        # To simplify implementation, please initialize both self.lin_dst\n",
        "        # and self.lin_src out_features to out_channels\n",
        "        self.lin_dst = None\n",
        "        self.lin_src = None\n",
        "\n",
        "        self.lin_update = None\n",
        "\n",
        "        ############# Your code here #############\n",
        "        ## (~3 lines of code)\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "\n",
        "    def forward(self, node_feature_src, node_feature_dst, edge_index, size=None, res_n_id=None):\n",
        "        ############# Your code here #############\n",
        "        ## (~1 line of code)\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "        return\n",
        "\n",
        "    def message_and_aggregate(self, edge_index, node_feature_src):\n",
        "\n",
        "        ############# Your code here #############\n",
        "        ## (~1 line of code)\n",
        "        ## Note:\n",
        "        ## 1. Different from what we implemented in Colab 3, we use message_and_aggregate\n",
        "        ## to replace the message and aggregate. The benefit is that we can avoid\n",
        "        ## materializing x_i and x_j, and make the implementation more efficient.\n",
        "        ## 2. To implement efficiently, following PyG documentation is helpful:\n",
        "        ## https://pytorch-geometric.readthedocs.io/en/latest/notes/sparse_tensor.html\n",
        "        ## 3. Here edge_index is torch_sparse SparseTensor.\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "\n",
        "        return\n",
        "\n",
        "    def update(self, aggr_out, node_feature_dst, res_n_id):\n",
        "\n",
        "        ############# Your code here #############\n",
        "        ## (~4 lines of code)\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "\n",
        "        return aggr_out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKq8ScTiJthn"
      },
      "source": [
        "### 异构 GNN 包装器\n",
        "\n",
        "在给每一种消息类型都实现 GNN 层之后，我们需要将消息关于类型聚合起来，这里我们将实现两种消息级别的聚合。\n",
        "\n",
        "第一种就是简单的平均聚合：\n",
        "\n",
        "$$\n",
        "\\begin{equation*}\n",
        "h_v^{(l)} = \\frac{1}{M}\\sum_{m=1}^{M}h_v^{(l)[m]}\n",
        "\\end{equation*}\n",
        "$$\n",
        "\n",
        "这里的结点 $v$ 的类型是 $d$，$M$ 是目标结点类型为 $d$ 的消息类型总数。\n",
        "\n",
        "另一种就是 **HAN** 中（[Wang et al. (2019)](https://arxiv.org/abs/1903.07293)）引入的语义级别注意力机制。与之前直接平均不一样，我们使用注意力机制来学习哪一种消息类型结果会重要，然后再聚合所有的消息类型。下列等式就是语义级别的注意力机制：\n",
        "\n",
        "$$\n",
        "\\begin{equation*}\n",
        "e_{m} = \\frac{1}{|V_{d}|} \\sum_{v \\in V_{d}} q_\\text{attn}^T \\cdot \\tanh \\left( W_\\text{attn}^{(l)} \\cdot h_v^{(l)[m]} + b \\right)\n",
        "\\end{equation*}\n",
        "$$\n",
        "\n",
        "其中 $m$ 指消息类型，$d$ 指目标结点类型。然后我们就可以计算注意力机制，并更新 $h_v^{(l)}$：\n",
        "\n",
        "$$\n",
        "\\begin{equation*}\n",
        "\\alpha_{m} = \\frac{\\exp(e_{m})}{\\sum_{m=1}^M \\exp(e_{m})}\n",
        "\\end{equation*}\n",
        "$$\n",
        "\n",
        "$$\n",
        "\\begin{equation*}\n",
        "h_v^{(l)} = \\sum_{m=1}^{M} \\alpha_{m} \\cdot h_v^{(l)[m]}\n",
        "\\end{equation*}\n",
        "$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0_bun02xJwFm"
      },
      "outputs": [],
      "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",
        "        self.attn_proj = None\n",
        "\n",
        "        if self.aggr == \"attn\":\n",
        "            ############# Your code here #############\n",
        "            ## (~1 line of code)\n",
        "            ## Note:\n",
        "            ## 1. Initialize self.attn_proj here.\n",
        "            ## 2. You should use nn.Sequential for self.attn_proj\n",
        "            ## 3. nn.Linear and nn.Tanh are useful.\n",
        "            ## 4. You can create a vector parameter by using:\n",
        "            ## nn.Linear(some_size, 1, bias=False)\n",
        "            ## 5. The first linear layer should have out_features as args['attn_size']\n",
        "            ## 6. You can assume we only have one \"head\" for the attention.\n",
        "            ## 7. We recommend you to implement the mean aggregation first. After \n",
        "            ## the mean aggregation works well in the training, then you can \n",
        "            ## implement this part.\n",
        "            pass\n",
        "\n",
        "            ##########################################\n",
        "    \n",
        "    def reset_parameters(self):\n",
        "        super(HeteroGNNWrapperConv, 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",
        "        # TODO: Implement this function that aggregates all message type results.\n",
        "        # Here, xs is a list of tensors (embeddings) with respect to message \n",
        "        # type aggregation results.\n",
        "\n",
        "        if self.aggr == \"mean\":\n",
        "\n",
        "            ############# Your code here #############\n",
        "            ## (~2 lines of code)\n",
        "            pass\n",
        "            ##########################################\n",
        "\n",
        "        elif self.aggr == \"attn\":\n",
        "\n",
        "            ############# Your code here #############\n",
        "            ## (~10 lines of code)\n",
        "            ## Note:\n",
        "            ## 1. Store the value of attention alpha (as a numpy array) to self.alpha,\n",
        "            ## which has the shape (len(xs), ) self.alpha will be not be used \n",
        "            ## to backpropagate etc. in the model. We will use it to see how much \n",
        "            ## attention the layer pays on different message types.\n",
        "            ## 2. torch.softmax and torch.cat are useful.\n",
        "            ## 3. You might need to reshape the tensors by using the \n",
        "            ## `view()` function https://pytorch.org/docs/stable/tensor_view.html\n",
        "            pass\n",
        "            ##########################################"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tn_pnCOKJw-d"
      },
      "source": [
        "### 初始化异构 GNN 层\n",
        "\n",
        "现在我们来初始化异构 GNN 层。与同构图不同，异构图有些难。\n",
        "\n",
        "总体而言，我们需要创建一个字典的 `HeteroGNNConv` 层，键为消息类型。\n",
        "\n",
        "- 想要获取所有的消息类型？`deepsnap.hetero_graph.HeteroGraph.message_types`会有帮助。\n",
        "- 如果我们初始化的**是第一个**卷积层，我们得要知道每种结点类型的特征维度。使用 `deepsnap.hetero_graph.HeteroGraph.num_node_features(node_type)` 能获取 `node_type` 类型结点的特征维度。\n",
        "- 如果我们初始化的**不是第一个**卷积层，所有结点类型都由一样的嵌入大小 `hidden_size`。\n",
        "- 无论初始化的是第一个卷积层与否，我们都将 `HeteroGNNConv` 的 `out_channels` 设置成 `hidden_size`。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hSBImHClJzf4"
      },
      "outputs": [],
      "source": [
        "def generate_convs(hetero_graph, conv, hidden_size, first_layer=False):\n",
        "    # TODO: Implement this function that returns a dictionary of `HeteroGNNConv` \n",
        "    # layers where the keys are message types. `hetero_graph` is deepsnap `HeteroGraph`\n",
        "    # object and the `conv` is the `HeteroGNNConv`.\n",
        "\n",
        "    convs = {}\n",
        "\n",
        "    ############# Your code here #############\n",
        "    ## (~9 lines of code)\n",
        "\n",
        "\n",
        "    ##########################################\n",
        "    \n",
        "    return convs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U39dX8EpJ3FG"
      },
      "source": [
        "### `HeteroGNN`\n",
        "\n",
        "现在我们来写一个只有两层 `HeteroGNNWrapperConv` 的简单 `HeteroGNN` 模型。\n",
        "\n",
        "`HeteroGNN` 中的 `forward` 方法会按照如下运行：\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",
      "execution_count": null,
      "metadata": {
        "id": "rplknA8aJ6J5"
      },
      "outputs": [],
      "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.convs1 = None\n",
        "        self.convs2 = None\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",
        "        ############# Your code here #############\n",
        "        ## (~10 lines of code)\n",
        "        ## Note:\n",
        "        ## 1. For self.convs1 and self.convs2, call generate_convs at first and then\n",
        "        ## pass the returned dictionary of `HeteroGNNConv` to `HeteroGNNWrapperConv`.\n",
        "        ## 2. For self.bns, self.relus and self.post_mps, the keys are node_types.\n",
        "        ## `deepsnap.hetero_graph.HeteroGraph.node_types` will be helpful.\n",
        "        ## 3. Initialize all batchnorms to torch.nn.BatchNorm1d(hidden_size, eps=1.0).\n",
        "        ## 4. Initialize all relus to nn.LeakyReLU().\n",
        "        ## 5. For self.post_mps, each value in the ModuleDict is a linear layer \n",
        "        ## where the `out_features` is the number of classes for that node type.\n",
        "        ## `deepsnap.hetero_graph.HeteroGraph.num_node_labels(node_type)` will be\n",
        "        ## useful.\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "\n",
        "    def forward(self, node_feature, edge_index):\n",
        "        # TODO: Implement the forward function. Notice that `node_feature` is \n",
        "        # a dictionary of tensors where keys are node types and values are \n",
        "        # corresponding feature tensors. The `edge_index` is a dictionary of \n",
        "        # tensors where keys are message types and values are corresponding\n",
        "        # edge index tensors (with respect to each message type).\n",
        "\n",
        "        x = node_feature\n",
        "\n",
        "        ############# Your code here #############\n",
        "        ## (~7 lines of code)\n",
        "        ## Note:\n",
        "        ## 1. `deepsnap.hetero_gnn.forward_op` can be helpful.\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "        \n",
        "        return x\n",
        "\n",
        "    def loss(self, preds, y, indices):\n",
        "        \n",
        "        loss = 0\n",
        "        loss_func = F.cross_entropy\n",
        "\n",
        "        ############# Your code here #############\n",
        "        ## (~3 lines of code)\n",
        "        ## Note:\n",
        "        ## 1. For each node type in preds, accumulate computed loss to `loss`\n",
        "        ## 2. Loss need to be computed with respect to the given index\n",
        "\n",
        "\n",
        "        ##########################################\n",
        "\n",
        "        return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9e7q_hUJ8zB"
      },
      "source": [
        "### 训练并测试\n",
        "\n",
        "现在我们给你提供了用来训练和测试的函数，你只需要实现一行代码。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CI5Hl_5TJ_YL"
      },
      "outputs": [],
      "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 = None\n",
        "\n",
        "    ############# Your code here #############\n",
        "    ## Note:\n",
        "    ## 1. `deepsnap.hetero_graph.HeteroGraph.node_label` is useful\n",
        "    ## 2. Compute the loss here\n",
        "\n",
        "\n",
        "    ##########################################\n",
        "\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DpNz9B5AKBUU"
      },
      "outputs": [],
      "source": [
        "# Please do not change the following parameters\n",
        "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",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pRHbWD4hKED8"
      },
      "source": [
        "### 数据集和预处理\n",
        "\n",
        "下一步，我们要加载数据集并创建张量后端（不使用 `NetworkX` 图）`deepsnap.hetero_graph.HeteroGraph` 对象。\n",
        "\n",
        "然后使用 **HAN**（[Wang et al. (2019)](https://arxiv.org/abs/1903.07293)）中提出的任务，即使用 `ACM(3025)` 数据集进行结点分类任务，我们的数据集是抽取自 [DGL](https://www.dgl.ai/) 的 [ACM.mat](https://data.dgl.ai/dataset/ACM.mat)。\n",
        "\n",
        "最初的 ACM 数据集有三种结点类型和两个边（关系）类型。简单起见，我们使用一种如下所示的结点类型和两种边类型的异构图。这意味着在我们的异构图中，我们只有一种结点类型（$\\text{paper}$）和两种消息类型（$(\\text{paper}\\xrightarrow{\\text{author}}\\text{paper})$ 和 $(\\text{paper}\\xrightarrow{\\text{subject}}\\text{paper})$）。\n",
        "\n",
        "<br/>\n",
        "<center>\n",
        "<img src=\"assets/cs224w-acm.png\"/>\n",
        "</center>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lJy03_IsKGh6"
      },
      "outputs": [],
      "source": [
        "print(f\"Device: {args['device']}\")\n",
        "\n",
        "# Load the data\n",
        "data = torch.load(\"dataset/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])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0HplV9hKMkc"
      },
      "source": [
        "### 训练平均聚合"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZgwfyzLbKOUw"
      },
      "outputs": [],
      "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",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rtkKBI_nKS1T"
      },
      "source": [
        "#### 问题 2.1: 使用平均聚合时，你用 `best_model` 在测试集上能获得的最大 **micro** F1 值是多少？"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AIvw51jMKTvn"
      },
      "source": [
        "#### 问题 2.2: 使用平均聚合时，你用 `best_model` 在测试集上能获得的最大 **macro** F1 值是多少？"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JBiYvwcuKd0z"
      },
      "source": [
        "### 训练注意力聚合"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e6na5zyQKfvi"
      },
      "outputs": [],
      "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",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mtAhFLQQKgbl"
      },
      "source": [
        "#### 问题 2.3: 使用注意力聚合时，你用 `best_model` 在测试集上能获得的最大 **micro** F1 值是多少？"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_cnsMGbsqJG_"
      },
      "source": [
        "#### 问题 2.4: 使用注意力聚合时，你用 `best_model` 在测试集上能获得的最大 **macro** F1 值是多少？"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQgx5y4UqMHH"
      },
      "source": [
        "### 关于消息类型的注意力机制\n",
        "\n",
        "通过消息类型等级的注意力机制，我们知道消息种类比来自消息来自哪一层更重要。\n",
        "\n",
        "下面我们打印并说明了每一层中不同消息类型的注意力情况。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hvK58gijqN_C"
      },
      "outputs": [],
      "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}\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "CS224W - Colab 4.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3.9.12 ('cs224w')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "3d401f214986802b882e8a583f41ee2435f6c7e7677ce5b0132581f351a1fd17"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}