# **CS224W - Colab 4**

In this Colab, we will shift our focus from homogenous graphs to heterogeneous graphs. Heterogeneous graphs extend the traditional homogenous graphs that we have been working with by incorperating different node and edge types. This additional information allows us to extend the graph neural nework models that we have worked with before. Namely, we can apply heterogenous message passing, where different message types now exist between different node and edge type relationships.

In this notebook, we will first learn how to transform NetworkX graphs into DeepSNAP representations. Then, we will dive deeper into how DeepSNAP stores and represents heterogeneous graphs as PyTorch Tensors.

Lastly, we will build our own heterogenous graph neural netowrk models using PyTorch Geometric and DeepSNAP. We will then apply our models for a node property prediction task; specifically, we will evaluate these models on the heterogeneous ACM node prediction dataset.

**Note**: Make sure to **sequentially run all the cells in each section**, so that the intermediate variables / packages will carry over to the next cell

Have fun and good luck on Colab 4 :)

# Device
You might need to use GPU for this Colab.

Please click `Runtime` and then `Change runtime type`. Then set the `hardware accelerator` to **GPU**.

# Installation

```python
# Install torch geometric
import os
if 'IS_GRADESCOPE_ENV' not in os.environ:
  !pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.9.0+cu111.html
  !pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.9.0+cu111.html
  !pip install torch-geometric
  !pip install -q git+https://github.com/snap-stanford/deepsnap.git
  !pip install -U -q PyDrive
```

    Looking in links: https://pytorch-geometric.com/whl/torch-1.9.0+cu111.html
    Requirement already satisfied: torch-scatter in /usr/local/lib/python3.7/dist-packages (2.0.9)
    Looking in links: https://pytorch-geometric.com/whl/torch-1.9.0+cu111.html
    Requirement already satisfied: torch-sparse in /usr/local/lib/python3.7/dist-packages (0.6.12)
    Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-sparse) (1.4.1)
    Requirement already satisfied: numpy>=1.13.3 in /usr/local/lib/python3.7/dist-packages (from scipy->torch-sparse) (1.19.5)
    Requirement already satisfied: torch-geometric in /usr/local/lib/python3.7/dist-packages (2.0.2)
    Requirement already satisfied: rdflib in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (6.0.2)
    Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.19.5)
    Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (0.22.2.post1)
    Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (4.62.3)
    Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.4.1)
    Requirement already satisfied: yacs in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (0.1.8)
    Requirement already satisfied: pyparsing in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.4.7)
    Requirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (3.13)
    Requirement already satisfied: jinja2 in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.11.3)
    Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (0.4)
    Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.1.5)
    Requirement already satisfied: networkx in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.6.3)
    Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.23.0)
    Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from jinja2->torch-geometric) (2.0.1)
    Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->torch-geometric) (2018.9)
    Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas->torch-geometric) (2.8.2)
    Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas->torch-geometric) (1.15.0)
    Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (57.4.0)
    Requirement already satisfied: isodate in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (0.6.0)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2.10)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (3.0.4)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2021.5.30)
    Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (1.24.3)
    Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->torch-geometric) (1.0.1)

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  !nvcc --version
  !python -c "import torch; print(torch.version.cuda)"
```

    nvcc: NVIDIA (R) Cuda compiler driver
    Copyright (c) 2005-2020 NVIDIA Corporation
    Built on Mon_Oct_12_20:09:46_PDT_2020
    Cuda compilation tools, release 11.1, V11.1.105
    Build cuda_11.1.TC455_06.29190527_0
    11.1

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  import torch
  import torch_geometric
  print(torch.__version__)
  print(torch_geometric.__version__)
```

    1.9.0+cu111
    2.0.2

# 1) DeepSNAP Heterogeneous Graph

First, we will explore how to transform a NetworkX graph into the format supported by DeepSNAP.

DeepSNAP extends its traditional graph representation to include heterogeneous graphs by including the following graph property features:
* `node_feature`: The feature of each node (`torch.tensor`)
* `edge_feature`: The feature of each edge (`torch.tensor`)
* `node_label`: The label of each node (`int`)
* `node_type`: The node type of each node (`string`)
* `edge_type`: The edge type of each edge (`string`)

where the key **new** features we add are `node_type` and `edge_type`, which enables us to perform heterogenous message passing.

For this first question we will work with the familiar [karate club graph](https://networkx.github.io/documentation/stable/auto_examples/graph/plot_karate_club.html) seen in Colab 1. To start, since each node in the graph belongs to one of two clubs (club "Mr. Hi" or club "Officer"), we will treat the club as the `node_type`. The code below demonstrates how to differentiate the nodes in the NetworkX graph.

```python
import networkx as nx
from networkx.algorithms.community import greedy_modularity_communities
import matplotlib.pyplot as plt
import copy

if 'IS_GRADESCOPE_ENV' not in os.environ:
  from pylab import show
  G = nx.karate_club_graph()
  community_map = {}
  for node in G.nodes(data=True):
    if node[1]["club"] == "Mr. Hi":
      community_map[node[0]] = 0
    else:
      community_map[node[0]] = 1
  node_color = []
  color_map = {0: 0, 1: 1}
  node_color = [color_map[community_map[node]] for node in G.nodes()]
  pos = nx.spring_layout(G)
  plt.figure(figsize=(7, 7))
  nx.draw(G, pos=pos, cmap=plt.get_cmap('coolwarm'), node_color=node_color)
  show()
```

### Question 1.1: Assigning Node Type and Node Features

Using the `community_map` dictionary and graph `G` from above, add node attributes `node_type` and `node_label` to the graph G. Namely, for `node_type` assign nodes in the "Mr. Hi" club to a node type `n0` and nodes in club "Officer" a node type `n1`. Note: the node type should be a `string` property.

Then for `node_label`, assign nodes in "Mr. Hi" club to a `node_label` `0` and nodes in club "Officer" a `node_label` of `1`.

Lastly, assign every node the *tensor* feature vector $[1, 1, 1, 1, 1]$.

**Hint**: Look at the NetworkX function `nx.classes.function.set_node_attributes`.

**Note**: This question is not specifically graded but is important for later questions.

```python
import torch

def assign_node_types(G, community_map):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and community map assignment (mapping node id --> 0/1 label)
  # and adds 'node_type' as a node_attribute in G.

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`
  ## 2. Look above for the two node type values!
  node_types = {node: 'n'+ str(community_map[node]) for node in G.nodes()}
  nx.set_node_attributes(G, node_types, 'node_type')

  #########################################

def assign_node_labels(G, community_map):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and community map assignment (mapping node id --> 0/1 label)
  # and adds 'node_label' as a node_attribute in G.

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`
  nx.set_node_attributes(G, community_map, 'node_label')

  #########################################

def assign_node_features(G):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and adds 'node_feature' as a node_attribute in G. Each node
  # in the graph has the same feature vector - a torchtensor with
  # data [1., 1., 1., 1., 1.]

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`
  nx.set_node_attributes(G, torch.Tensor([1., 1., 1., 1., 1.]), 'node_feature')

  #########################################

if 'IS_GRADESCOPE_ENV' not in os.environ:
  assign_node_types(G, community_map)
  assign_node_labels(G, community_map)
  assign_node_features(G)

  # Explore node properties for the node with id: 20
  node_id = 20
  print (f"Node {node_id} has properties:", G.nodes(data=True)[node_id])
```

    Node 20 has properties: {'club': 'Officer', 'node_type': 'n1', 'node_label': 1, 'node_feature': tensor([1., 1., 1., 1., 1.])}

### Question 1.2: Assigning Edge Types

Next, we will assign three different `edge_types`:
* Edges within club "Mr. Hi": `e0`
* Edges within club "Officer": `e1`
* Edges between the two clubs: `e2`

**Hint**: Use the `community_map` from before and `nx.classes.function.set_edge_attributes`

```python
def assign_edge_types(G, community_map):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and community map assignment (mapping node id --> 0/1 label)
  # and adds 'edge_type' as a edge_attribute in G.

  ############# Your code here ############
  ## (~5 line of code)
  ## Note
  ## 1. Create an edge assignment dict following rules above
  ## 2. Look up NetworkX `nx.classes.function.set_edge_attributes`

  edge_map = {'00': 'e0', '01': 'e2', '10': 'e2', '11': 'e1'}
  edge_types = {edge: edge_map[str(community_map[edge[0]]) + str(community_map[edge[1]])] for edge in G.edges}
  nx.set_edge_attributes(G, edge_types, 'edge_type')
  #########################################

if 'IS_GRADESCOPE_ENV' not in os.environ:
  assign_edge_types(G, community_map)

  # Explore edge properties for a sampled edge and check the corresponding
  # node types
  edge_idx = 15
  n1 = 0
  n2 = 31
  edge = list(G.edges(data=True))[edge_idx]
  print (f"Edge ({edge[0]}, {edge[1]}) has properties:", edge[2])
  print (f"Node {n1} has properties:", G.nodes(data=True)[n1])
  print (f"Node {n2} has properties:", G.nodes(data=True)[n2])
```

    Edge (0, 31) has properties: {'edge_type': 'e2'}
    Node 0 has properties: {'club': 'Mr. Hi', 'node_type': 'n0', 'node_label': 0, 'node_feature': tensor([1., 1., 1., 1., 1.])}
    Node 31 has properties: {'club': 'Officer', 'node_type': 'n1', 'node_label': 1, 'node_feature': tensor([1., 1., 1., 1., 1.])}

## Heterogeneous Graph Visualization

Now we can visualize the Heterogeneous Graph we have generated.

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  edge_color = {}
  for edge in G.edges():
    n1, n2 = edge
    edge_color[edge] = community_map[n1] if community_map[n1] == community_map[n2] else 2
    if community_map[n1] == community_map[n2] and community_map[n1] == 0:
      edge_color[edge] = 'blue'
    elif community_map[n1] == community_map[n2] and community_map[n1] == 1:
      edge_color[edge] = 'red'
    else:
      edge_color[edge] = 'green'

  G_orig = copy.deepcopy(G)
  nx.classes.function.set_edge_attributes(G, edge_color, name='color')
  colors = nx.get_edge_attributes(G,'color').values()
  labels = nx.get_node_attributes(G, 'node_type')
  plt.figure(figsize=(8, 8))
  nx.draw(G, pos=pos, cmap=plt.get_cmap('coolwarm'), node_color=node_color, edge_color=colors, labels=labels, font_color='white')
  show()
```

where we differentiate edges within each clubs (2 types) and edges between the two clubs (1 type). Different types of nodes and edges are visualized in different colors. The NetworkX object `G` in following code can be transformed into `deepsnap.hetero_graph.HeteroGraph` directly.

## Transforming to DeepSNAP representation

We will now work through transforming the NetworkX object `G` into a `deepsnap.hetero_graph.HeteroGraph`.

```python
from deepsnap.hetero_graph import HeteroGraph

if 'IS_GRADESCOPE_ENV' not in os.environ:
  hete = HeteroGraph(G_orig)
```

## Question 1.3: How many nodes are of each type (10 Points)

```python
def get_nodes_per_type(hete):
  # TODO: Implement a function that takes a DeepSNAP dataset object
  # and return the number of nodes per `node_type`.

  num_nodes_n0 = 0
  num_nodes_n1 = 0

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Colab autocomplete functionality might be useful.
  num_nodes_n0 = len(hete.node_type['n0'])
  num_nodes_n1 = len(hete.node_type['n1'])

  #########################################

  return num_nodes_n0, num_nodes_n1

if 'IS_GRADESCOPE_ENV' not in os.environ:
  num_nodes_n0, num_nodes_n1 = get_nodes_per_type(hete)
  print("Node type n0 has {} nodes".format(num_nodes_n0))
  print("Node type n1 has {} nodes".format(num_nodes_n1))
```

    Node type n0 has 17 nodes
    Node type n1 has 17 nodes

## Question 1.4: Message Types - How many edges are of each message type (10 Points)

When working with heterogenous graphs, as we have discussed before, we now work with heterogenous message types (i.e. different message types for the different `node_type` and `edge_type` combinations). For example, an edge of type `e0` connecting two nodes in club "Mr. HI" would have a message type of (`n0`, `e0`, `n0`). In this problem we will analyze how many edges in our graph are of each message type.

**Hint**: If you want to learn more about what the different message types are try the call `hete.message_types`

```python
def get_num_message_edges(hete):
  # TODO: Implement this function that takes a DeepSNAP dataset object
  # and return the number of edges for each message type.
  # You should return a list of tuples as
  # (message_type, num_edge)

  message_type_edges = []

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Colab autocomplete functionality might be useful.
  for message_type in hete.message_types:
    message_type_edges.append((message_type, len(hete.edge_type[message_type])))

  #########################################

  return message_type_edges

if 'IS_GRADESCOPE_ENV' not in os.environ:
  message_type_edges = get_num_message_edges(hete)
  for (message_type, num_edges) in message_type_edges:
    print("Message type {} has {} edges".format(message_type, num_edges))
```

    Message type ('n0', 'e0', 'n0') has 35 edges
    Message type ('n0', 'e2', 'n1') has 11 edges
    Message type ('n1', 'e1', 'n1') has 32 edges

## Question 1.5: Dataset Splitting - How many nodes are in each dataset split? (10 Points)

DeepSNAP has built in Dataset creation and splitting methods for heterogeneous graphs. Here we will create train, validation, and test datasets for a node prediction task and inspect the resulting subgraphs. Specifically, write a function that computes the number of nodes with a known label in each dataset split.

```python
from deepsnap.dataset import GraphDataset

def compute_dataset_split_counts(datasets):
  # TODO: Implement a function that takes a dict of datasets in the form
  # {'train': dataset_train, 'val': dataset_val, 'test': dataset_test}
  # and returns a dict mapping dataset names to the number of labeled
  # nodes used for supervision in that respective dataset.

  data_set_splits = {}

  ############# Your code here ############
  ## (~3 line of code)
  ## Note
  ## 1. The DeepSNAP `node_label_index` dictionary will be helpful.
  ## 2. Remember to count both node_types
  ## 3. Remember each dataset only has one graph that we need to access
  ##    (i.e. dataset[0])
  for key in datasets.keys():
    num_nodes_n0 = datasets[key][0].node_label_index['n0'].shape[0]
    num_nodes_n1 = datasets[key][0].node_label_index['n1'].shape[0]
    data_set_splits[key] = num_nodes_n0 + num_nodes_n1

  #########################################

  return data_set_splits

if 'IS_GRADESCOPE_ENV' not in os.environ:
  dataset = GraphDataset([hete], task='node')
  # Splitting the dataset
  dataset_train, dataset_val, dataset_test = dataset.split(transductive=True, split_ratio=[0.4, 0.3, 0.3])
  datasets = {'train': dataset_train, 'val': dataset_val, 'test': dataset_test}

  data_set_splits = compute_dataset_split_counts(datasets)
  for dataset_name, num_nodes in data_set_splits.items():
    print("{} dataset has {} nodes".format(dataset_name, num_nodes))
```

    train dataset has 12 nodes
    val dataset has 10 nodes
    test dataset has 12 nodes

## DeepSNAP Dataset Visualization

We can now visualize the different nodes and edges used in each graph dataset split.

```python
from deepsnap.dataset import GraphDataset

if 'IS_GRADESCOPE_ENV' not in os.environ:
  dataset = GraphDataset([hete], task='node')
  # Splitting the dataset
  dataset_train, dataset_val, dataset_test = dataset.split(transductive=True, split_ratio=[0.4, 0.3, 0.3])
  titles = ['Train', 'Validation', 'Test']

  for i, dataset in enumerate([dataset_train, dataset_val, dataset_test]):
    n0 = hete._convert_to_graph_index(dataset[0].node_label_index['n0'], 'n0').tolist()
    n1 = hete._convert_to_graph_index(dataset[0].node_label_index['n1'], 'n1').tolist()

    plt.figure(figsize=(7, 7))
    plt.title(titles[i])
    nx.draw(G_orig, pos=pos, node_color="grey", edge_color=colors, labels=labels, font_color='white')
    nx.draw_networkx_nodes(G_orig.subgraph(n0), pos=pos, node_color="blue")
    nx.draw_networkx_nodes(G_orig.subgraph(n1), pos=pos, node_color="red")
    show()
```

# 2) Heterogeneous Graph Node Property Prediction

Now, we will use PyTorch Geometric and DeepSNAP to implement a GNN model for heterogeneous graph node property prediction (node classification). We will draw upon our understanding of heterogeneous graphs from lecture and previous work in implementing GNN layers using PyG (introduced in Colab 3).

First let's take a look at the general structure of a heterogeneous GNN layer by working through an example:

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)$. Note: during message passing we view each message as (src, relation, dst), where messages "flow" from src to dst node types. For example, during message passing, updating node type $b$ relies on two different message types $m_2$ and $m_3$.

When applying message passing in heterogenous graphs, we seperately apply message passing over each message type. Therefore, for the graph $G$, a heterogeneous GNN layer contains three seperate Heterogeneous Message Passing layers (`HeteroGNNConv` in this Colab), where each `HeteroGNNConv` layer performs message passing and aggregation with respect to *only one message type*. Since a message type is viewed as (src, relation, dst) and messages "flow" from src to dst, each `HeteroGNNConv` layer only computes embeddings for the *dst* nodes of a given message type. For example, the `HeteroGNNConv` layer for message type $m_2$ outputs updated embedding representations *only* for node's with type b.

---

An overview of the heterogeneous layer we will create is shown below:

where we highlight the following notation:

- $H_a^{(l)[m_1]}$ is the intermediate matrix of of node embeddings for node type $a$, generated by the $l$th `HeteroGNNConv` layer for message type $m_1$.
- $H_a^{(l)}$ is the matrix with current embeddings for nodes of type $a$ after the $l$th layer of our Heterogeneous GNN model. Note that these embeddings can rely on one or more intermediate `HeteroGNNConv` layer embeddings(i.e. $H_b^{(l)}$ combines $H_b^{(l)[m_2]}$ and $H_b^{(l)[m_3]}$).

Since each `HeteroGNNConv` is only applied over a single message type, we additionally define a Heterogeneous GNN Wrapper layer (`HeteroGNNWrapperConv`). This wrapper manages and combines the output of each `HeteroGNNConv` layer in order to generate the complete updated node embeddings for each node type in layer $l$ of our model. More specifically, the $l^{th}$ `HeteroGNNWrapperConv` layer takes as input the node embeddings computed for each message type and node type (e.g. $H_b^{(l)[m_2]}$ and $H_b^{(l)[m_3]}$) and aggregates across message types with the same $dst$ node type. The resulting output of the $l^{th}$ `HeteroGNNWrapperConv` layer is the updated embedding matrix $H_i^{(l)}$ for each node type i.

Continuing on our example above, to compute the node embeddings $H_b^{(l)}$ the wrapper layer aggregates output embeddings from the `HeteroGNNConv` layers associated with message types $m_2$ and $m_3$ (i.e. $H_b^{(l)[m_2]}$ and $H_b^{(l)[m_3]}$).

---

With the `HeteroGNNWrapperConv` module, we can now draw a "simplified" heterogeneous layer structure as follows:

---
**NOTE**:
As reference, it may be helpful to additionally read through PyG's introduciton to heterogeneous graph representations and buidling heterogeneous GNN models: https://pytorch-geometric.readthedocs.io/en/latest/notes/heterogeneous.html

Looking ahead, we recommend you implement the heterogeneous GNN model in following steps:

1. Implement `HeteroGNNConv`.
2. Implement **just** `mean` aggregation within `HeteroGNNWrapperConv`.
3. Implement `generate_convs`.
4. Implement the `HeteroGNN` model and the `train` function.
5. Train the model with `mean` aggregation and test your model to make sure your model has reasonable performance.
6. Once you are confident in your mean aggregation model, implement `attn` aggregation in `HeteroGNNWrapperConv`.
7. Train the model with `attn` aggregation and test your model to make sure your model has reasonable performance.

Note: The key point of advice is to work completely through implementing the mean aggregation heterogeneous GNN model before diving into the more difficult attention based model.

## Setup

```python
import copy
import torch
import deepsnap
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch_geometric.nn as pyg_nn

from sklearn.metrics import f1_score
from deepsnap.hetero_gnn import forward_op
from deepsnap.hetero_graph import HeteroGraph
from torch_sparse import SparseTensor, matmul
```

## Dataset

You need to login to your Google account and enter the verification code below.

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  from pydrive.auth import GoogleAuth
  from pydrive.drive import GoogleDrive
  from google.colab import auth
  from oauth2client.client import GoogleCredentials

  # Authenticate and create the PyDrive client
  auth.authenticate_user()
  gauth = GoogleAuth()
  gauth.credentials = GoogleCredentials.get_application_default()
  drive = GoogleDrive(gauth)
```

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  id='1ivlxd6lJMcZ9taS44TMGG72x2V1GeVvk'
  downloaded = drive.CreateFile({'id': id})
  downloaded.GetContentFile('acm.pkl')
```

## Implementing `HeteroGNNConv`

Now let's start working on our own implementation of the heterogeneous message passing layer (`HeteroGNNConv`)! Just as in Colab 3, we will implement the layer using PyTorch Geometric.

At a high level, the `HeteroGNNConv` layer is equivalent to the homogenous GNN layers we implemented in Colab 3, but now applied to an individual heterogeous message type. Moreover, our heterogeneous GNN layer draws directly from the **GraphSAGE** message passing model ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)).

We begin by defining the `HeteroGNNConv` layer with respect to message type $m$:

\begin{equation}
m =(s, r, d)
\end{equation}

where each message type is a tuple containing three elements: $s$ - the source node type, $r$ - the edge (relation) type, and $d$ - the destination node type.

The message passing update rule that we implement is very similar to that of GraphSAGE, except we now need to include the node types and the edge relation type. The update rule for message type $m$ is described below:

\begin{equation}
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)
\end{equation}

where we compute $h_v^{(l)[m]}$, the node embedding representation for node $v$ after `HeteroGNNConv` layer $l$ with respect message type $m$. Further unpacking the forumla we have:
- $W_s^{(l)[m]}$ - linear transformation matrix for the messages of neighboring source nodes of type $s$ along message type $m$.
- $W_d^{(l)[m]}$ - linear transformation matrix for the message from the node $v$ itself of type $d$.
- $W^{(l)[m]}$ - linear transformation matrix for the concatenated messages from neighboring node's and the central node.
- $h_u^{(l-1)}$ - the hidden embedding representation for node $u$ after the $(l-1)$th `HeteroGNNWrapperConv` layer. Note, that this embedding is not associated with a particular message type (see layer diagrams above).
- $N_{m}(v)$ - the set of neighbor source nodes $s$ for the node v that we are embedding along message type $m = (s, r, d)$.

**NOTE**: We emphasize that each weight matrix is associated with a specific message type $[m]$ and additionally, the weight matrices applied to node messages are differentiated by node type (i.e. $W_s$ and $W_d$).

Lastly, for simplicity, we use mean aggregations for $AGG$ where:

\begin{equation}
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)}
\end{equation}

```python
class HeteroGNNConv(pyg_nn.MessagePassing):
    def __init__(self, in_channels_src, in_channels_dst, out_channels):
        super(HeteroGNNConv, self).__init__(aggr="mean")

        self.in_channels_src = in_channels_src
        self.in_channels_dst = in_channels_dst
        self.out_channels = out_channels

        # To simplify implementation, please initialize both self.lin_dst
        # and self.lin_src out_features to out_channels
        self.lin_dst = None
        self.lin_src = None

        self.lin_update = None

        ############# Your code here #############
        ## (~3 lines of code)
        ## Note:
        ## 1. Initialize the 3 linear layers.
        ## 2. Think through the connection between the mathematical
        ##    definition of the update rule and torch linear layers!
        self.lin_dst = nn.Linear(self.in_channels_dst, self.out_channels)
        self.lin_src = nn.Linear(self.in_channels_src, self.out_channels)
        self.lin_update = nn.Linear(self.out_channels*2, self.out_channels)

        ##########################################

    def forward(
        self,
        node_feature_src,
        node_feature_dst,
        edge_index,
        size=None
    ):
        ############# Your code here #############
        ## (~1 line of code)
        ## Note:
        ## 1. Unlike Colab 3, we just need to call self.propagate with
        ## proper arguments.

        return self.propagate(
            edge_index,
            size=size,
            node_feature_src=node_feature_src,
            node_feature_dst=node_feature_dst,
        )
        ##########################################

    def message_and_aggregate(self, edge_index, node_feature_src):

        out = None
        ############# Your code here #############
        ## (~1 line of code)
        ## Note:
        ## 1. Different from what we implemented in Colab 3, we use message_and_aggregate
        ##    to combine the previously seperate message and aggregate functions.
        ##    The benefit is that we can avoid materializing x_i and x_j
        ##    to make the implementation more efficient.
        ## 2. To implement efficiently, refer to PyG documentation for message_and_aggregate
        ##    and sparse-matrix multiplication:
        ##    https://pytorch-geometric.readthedocs.io/en/latest/notes/sparse_tensor.html
        ## 3. Here edge_index is torch_sparse SparseTensor. Although interesting, you
        ##    do not need to deeply understand SparseTensor represenations!
        ## 4. Conceptually, think through how the message passing and aggregation
        ##    expressed mathematically can be expressed through matrix multiplication.

        out = matmul(edge_index, node_feature_src, reduce='mean')
        ##########################################

        return out

    def update(self, aggr_out, node_feature_dst):

        ############# Your code here #############
        ## (~4 lines of code)
        ## Note:
        ## 1. The update function is called after message_and_aggregate
        ## 2. Think through the one-one connection between the mathematical update
        ##    rule and the 3 linear layers defined in the constructor.

        node_feature_dst = self.lin_dst(node_feature_dst)
        aggr_out = self.lin_src(aggr_out)
        combined = torch.cat((node_feature_dst, aggr_out), dim=-1)
        aggr_out = self.lin_update(combined)
        ##########################################

        return aggr_out
```

## Heterogeneous GNN Wrapper Layer

After implementing the `HeteroGNNConv` layer for each message type, we need to manage and aggregate the node embedding results (with respect to each message types). Here we will implement two types of message type level aggregation.

The first one is simply mean aggregation over message types:

\begin{equation}
h_v^{(l)} = \frac{1}{M}\sum_{m=1}^{M}h_v^{(l)[m]}
\end{equation}

where node $v$ has node type $d$ and we sum over the $M$ message types that have destination node type $d$. From our original example, for a node v of type $b$ we aggregate v's `HeteroGNNConv` embeddings for message types $m_2$ and $m_3$ (i.e. $h_v^{(l)[m_2]}$ and $h_v^{(l)[m_3]}$).

The second method we implement 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 use attention to learn which message type result is more important, then aggregate across all the message types. Below are the equations for semantic level attention:

\begin{equation}
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)
\end{equation}

where $m$ is the message type and $d$ refers to the destination node type for that message ($m = (s, r, d)$). Additionally, $V_{d}$ refers to the set of nodes v with type $d$. Lastly, the unormalized attention weight $e_m$ is a scaler computed for each message type $m$.

Next, we can compute the normalized attention weights and update $h_v^{(l)}$:

\begin{equation}
\alpha_{m} = \frac{\exp(e_{m})}{\sum_{m=1}^M \exp(e_{m})}
\end{equation}

\begin{equation}
h_v^{(l)} = \sum_{m=1}^{M} \alpha_{m} \cdot h_v^{(l)[m]}
\end{equation}

, where we emphasize that $M$ here is the number of message types associated with the destination node type $d$.

**Note**: The implementation of the attention aggregation is tricky and nuanced. We strongly recommend working carefully through the math equations to undersatnd exactly what each notation refers to and how all the pieces fit together. If you can, try to connect the math to our original example, focusing on node type $b$, which depends on two different message types!

```python
class HeteroGNNWrapperConv(deepsnap.hetero_gnn.HeteroConv):
    def __init__(self, convs, args, aggr=None):
        """
        args = {
            'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
            'hidden_size': 64,
            'epochs': 100,
            'weight_decay': 1e-5,
            'lr': 0.003,
            'attn_size': 32,
        }
        """
        super(HeteroGNNWrapperConv, self).__init__(convs, None)
        self.aggr = aggr

        # Map the index and message type
        self.mapping = {}

        # A numpy array that stores the final attention probability
        self.alpha = None

        self.attn_proj = None

        if self.aggr == "attn":
            ############# Your code here #############
            ## (~1 line of code)
            ## Note:
            ## 1. Initialize self.attn_proj, where self.attn_proj should include
            ##    two linear layers. Note, make sure you understand
            ##    which part of the equation self.attn_proj captures.
            ## 2. You should use nn.Sequential for self.attn_proj
            ## 3. nn.Linear and nn.Tanh are useful.
            ## 4. You can model a weight vector (rather than matrix) by using:
            ##    nn.Linear(some_size, 1, bias=False).
            ## 5. The first linear layer should have out_features as args['attn_size']
            ## 6. You can assume we only have one "head" for the attention.
            ## 7. We recommend you to implement the mean aggregation first. After
            ##    the mean aggregation works well in the training, then you can
            ##    implement this part.

            self.attn_proj = nn.Sequential(
                nn.Linear(args['hidden_size'], args['attn_size']),
                nn.Tanh(),
                nn.Linear(args['attn_size'], 1, bias=False)
            )
            ##########################################

    def reset_parameters(self):
        super(HeteroGNNWrapperConv, self).reset_parameters()
        if self.aggr == "attn":
            for layer in self.attn_proj.children():
                layer.reset_parameters()

    def forward(self, node_features, edge_indices):
        message_type_emb = {}
        for message_key, message_type in edge_indices.items():
            src_type, edge_type, dst_type = message_key
            node_feature_src = node_features[src_type]
            node_feature_dst = node_features[dst_type]
            edge_index = edge_indices[message_key]
            message_type_emb[message_key] = (
                self.convs[message_key](
                    node_feature_src,
                    node_feature_dst,
                    edge_index,
                )
            )
        node_emb = {dst: [] for _, _, dst in message_type_emb.keys()}
        mapping = {}
        for (src, edge_type, dst), item in message_type_emb.items():
            mapping[len(node_emb[dst])] = (src, edge_type, dst)
            node_emb[dst].append(item)
        self.mapping = mapping
        for node_type, embs in node_emb.items():
            if len(embs) == 1:
                node_emb[node_type] = embs[0]
            else:
                node_emb[node_type] = self.aggregate(embs)
        return node_emb

    def aggregate(self, xs):
        # TODO: Implement this function that aggregates all message type results.
        # Here, xs is a list of tensors (embeddings) with respect to message
        # type aggregation results.

        # Useful dimensions from xs - particularly for `attn` aggregation
        N = xs[0].shape[0] # Number of nodes for that node type
        M = len(xs) # Number of message types for that node type

        if self.aggr == "mean":

            ############# Your code here #############
            ## (~2 lines of code)
            ## Note:
            ## 1. Explore the function parameter `xs`!

            # mean aggregation over message types

            # xs shape (M x N)
            h_v_lm = torch.stack(xs)                        # h_v_lm shape: M x N x hidden_size
            h_v_l = torch.mean(torch.stack(xs), dim=0)      # h_v_l shape: N x hidden_Size

            return h_v_l

            ##########################################

        elif self.aggr == "attn":
            ############# Your code here #############
            ## (~10 lines of code)
            ## Note:
            ## 1. Try to map out how the equations can be translated into code.
            ## 2. N and M defined above may be useful at least to understand.
            ## 3. Work first to compute the un-normalized attention weights e
            ##    for each message type - try to vectorize this!
            ## 4. torch.softmax and torch.cat are useful.
            ## 5. It might be useful to reshape and concatenate tensors using the
            ##    `view()` function https://pytorch.org/docs/stable/tensor_view.html
            ##    and `torch.cat()`https://pytorch.org/docs/stable/generated/torch.cat.html
            ## 6. Store the value of attention alpha (as a numpy array) to self.alpha,
            ##    which has the shape (len(xs), ) self.alpha will be not be used
            ##    to backpropagate etc. in the model. We will use it to see how much
            ##    attention the layer pays on different message types

            # semantic level attention

            # xs shape (M x N)
            h_v_lm = torch.stack(xs, dim=0)                 # h_v_lm shape: (M x N x hidden_size)

            e_m = self.attn_proj(h_v_lm)                    # e_m shape: (M x N x 1)
            e_m = torch.mean(e_m, dim=-1)                   # e_m shape: (M x N)

            self.alpha = torch.softmax(e_m, dim=0).detach() # alpha shape: (M x N)

            h_v_l = self.alpha.unsqueeze(-1) * h_v_lm       # h_v_l shape: (M x N x hidden_size)
            h_v_l = torch.sum(h_v_l, dim=0)                 # h_v_l shape: (N x hidden)

            return h_v_l
            ##########################################
```

## Initialize Heterogeneous GNN Layers

Now let's put it all together and initialize the Heterogeneous GNN Layers. Different from the homogeneous graph case, heterogeneous graphs can be a little bit complex.

In general, we need to create a dictionary of `HeteroGNNConv` layers where the keys are message types.

* To get all message types, `deepsnap.hetero_graph.HeteroGraph.message_types` is useful.
* If we are initializing the first conv layers, we need to get the feature dimension of each node type. Using `deepsnap.hetero_graph.HeteroGraph.num_node_features(node_type)` will return the node feature dimension of `node_type`. In this function, we will set each `HeteroGNNConv` `out_channels` to be `hidden_size`.
* If we are not initializing the first conv layers, all node types will have the smae embedding dimension `hidden_size` and we still set `HeteroGNNConv` `out_channels` to be `hidden_size` for simplicity.

```python
def generate_convs(hetero_graph, conv, hidden_size, first_layer=False):
    # TODO: Implement this function that returns a dictionary of `HeteroGNNConv`
    # layers where the keys are message types. `hetero_graph` is deepsnap `HeteroGraph`
    # object and the `conv` is the `HeteroGNNConv`.

    convs = {}

    ############# Your code here #############
    ## (~9 lines of code)
    ## Note:
    ## 1. See the hints above!
    ## 2. conv is of type `HeteroGNNConv`
    for message_type in hetero_graph.message_types:
      if first_layer:
        # each message is (src, rel, dst), so we take first and last
        num_node_feat_src = hetero_graph.num_node_features(message_type[0])
        num_node_feat_dst = hetero_graph.num_node_features(message_type[-1])
        convs[message_type] = conv(
            in_channels_src=num_node_feat_src,
            in_channels_dst=num_node_feat_dst,
            out_channels=hidden_size,
        )
      else:
        convs[message_type] = conv(
            in_channels_src=hidden_size,
            in_channels_dst=hidden_size,
            out_channels=hidden_size,
        )
    ##########################################

    return convs
```

## HeteroGNN

Now we will make a simple HeteroGNN model which contains only two `HeteroGNNWrapperConv` layers.

For the forward function in `HeteroGNN`, the model is going to be run as following:

$\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}$

```python
class HeteroGNN(torch.nn.Module):
    def __init__(self, hetero_graph, args, aggr="mean"):
        super(HeteroGNN, self).__init__()

        self.aggr = aggr
        self.hidden_size = args['hidden_size']

        self.convs1 = None
        self.convs2 = None

        self.bns1 = nn.ModuleDict()
        self.bns2 = nn.ModuleDict()
        self.relus1 = nn.ModuleDict()
        self.relus2 = nn.ModuleDict()
        self.post_mps = nn.ModuleDict()

        ############# Your code here #############
        ## (~10 lines of code)
        ## Note:
        ## 1. For self.convs1 and self.convs2, call generate_convs at first and then
        ##    pass the returned dictionary of `HeteroGNNConv` to `HeteroGNNWrapperConv`.
        ## 2. For self.bns, self.relus and self.post_mps, the keys are node_types.
        ##    `deepsnap.hetero_graph.HeteroGraph.node_types` will be helpful.
        ## 3. Initialize all batchnorms to torch.nn.BatchNorm1d(hidden_size, eps=1).
        ## 4. Initialize all relus to nn.LeakyReLU().
        ## 5. For self.post_mps, each value in the ModuleDict is a linear layer
        ##    where the `out_features` is the number of classes for that node type.
        ##    `deepsnap.hetero_graph.HeteroGraph.num_node_labels(node_type)` will be
        ##    useful.

        # Step 1
        convs1_dict = generate_convs(
            hetero_graph=hetero_graph,
            conv=HeteroGNNConv,
            hidden_size=self.hidden_size,
            first_layer=True,
        )
        self.convs1 = HeteroGNNWrapperConv(
            convs=convs1_dict,
            args=args,
            aggr=self.aggr
        )
        convs2_dict = generate_convs(
            hetero_graph=hetero_graph,
            conv=HeteroGNNConv,
            hidden_size=self.hidden_size,
            first_layer=False,
        )
        self.convs2 = HeteroGNNWrapperConv(
            convs=convs2_dict,
            args=args,
            aggr=self.aggr
        )

        # Step 2, 3, 4
        for node_type in hetero_graph.node_types:
          self.bns1[node_type] = torch.nn.BatchNorm1d(self.hidden_size, eps=1.)
          self.relus1[node_type] = torch.nn.LeakyReLU()
          self.bns2[node_type] = torch.nn.BatchNorm1d(self.hidden_size, eps=1.)
          self.relus2[node_type] = torch.nn.LeakyReLU()
          self.post_mps[node_type] = nn.Linear(
              self.hidden_size,
              hetero_graph.num_node_labels(node_type)
          )
        ##########################################

    def forward(self, node_feature, edge_index):
        # TODO: Implement the forward function. Notice that `node_feature` is
        # a dictionary of tensors where keys are node types and values are
        # corresponding feature tensors. The `edge_index` is a dictionary of
        # tensors where keys are message types and values are corresponding
        # edge index tensors (with respect to each message type).

        x = node_feature

        ############# Your code here #############
        ## (~7 lines of code)
        ## Note:
        ## 1. `deepsnap.hetero_gnn.forward_op` can be helpful for
        ##    the bn, relu, and post_mp ops.
        x = self.convs1(x, edge_index)
        x = forward_op(x, self.bns1)
        x = forward_op(x, self.relus1)

        x = self.convs2(x, edge_index)
        x = forward_op(x, self.bns2)
        x = forward_op(x, self.relus2)

        x = forward_op(x, self.post_mps)
        ##########################################

        return x

    def loss(self, preds, y, indices):

        loss = 0
        loss_func = F.cross_entropy

        ############# Your code here #############
        ## (~3 lines of code)
        ## Note:
        ## 1. For each node type in preds, accumulate computed loss to `loss`
        ## 2. Loss need to be computed with respect to the given index
        ## 3. preds is a dictionary of model predictions keyed by node_type.
        ## 4. indeces is a dictionary of labeled supervision nodes keyed
        ##    by node_type
        for node_type in preds:                             # 1.
          index = indices[node_type]                        # 4.
          pred = preds[node_type][index]                    # 3.
          loss += loss_func(pred, y[node_type][index])      # 2.
        ##########################################

        return loss
```

## Training and Testing

Here we provide you with the functions to train and test. You only need to implement one line of code here.

**Please do not modify other parts in `train` and `test` for grading purposes.**

```python
import pandas as pd

def train(model, optimizer, hetero_graph, train_idx):
    model.train()
    optimizer.zero_grad()
    preds = model(hetero_graph.node_feature, hetero_graph.edge_index)

    loss = None

    ############# Your code here #############
    ## Note:
    ## 1. Compute the loss here
    ## 2. `deepsnap.hetero_graph.HeteroGraph.node_label` is useful

    loss = model.loss(
        preds=preds,
        y=hetero_graph.node_label,
        indices=train_idx,
    )
    ##########################################

    loss.backward()
    optimizer.step()
    return loss.item()

def test(model, graph, indices, best_model=None, best_val=0, save_preds=False, agg_type=None):
    model.eval()
    accs = []
    for i, index in enumerate(indices):
        preds = model(graph.node_feature, graph.edge_index)
        num_node_types = 0
        micro = 0
        macro = 0
        for node_type in preds:
            idx = index[node_type]
            pred = preds[node_type][idx]
            pred = pred.max(1)[1]
            label_np = graph.node_label[node_type][idx].cpu().numpy()
            pred_np = pred.cpu().numpy()
            micro = f1_score(label_np, pred_np, average='micro')
            macro = f1_score(label_np, pred_np, average='macro')
            num_node_types += 1

        # Averaging f1 score might not make sense, but in our example we only
        # have one node type
        micro /= num_node_types
        macro /= num_node_types
        accs.append((micro, macro))

        # Only save the test set predictions and labels!
        if save_preds and i == 2:
          print ("Saving Heterogeneous Node Prediction Model Predictions with Agg:", agg_type)
          print()

          data = {}
          data['pred'] = pred_np
          data['label'] = label_np

          df = pd.DataFrame(data=data)
          # Save locally as csv
          df.to_csv('ACM-Node-' + agg_type + 'Agg.csv', sep=',', index=False)

    if accs[1][0] > best_val:
        best_val = accs[1][0]
        best_model = copy.deepcopy(model)
    return accs, best_model, best_val
```

```python
# Please do not change the following parameters
args = {
    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
    'hidden_size': 64,
    'epochs': 100,
    'weight_decay': 1e-5,
    'lr': 0.003,
    'attn_size': 32,
}
```

## Dataset and Preprocessing

In the next, we will load the data and create a tensor backend (without a NetworkX graph) `deepsnap.hetero_graph.HeteroGraph` object.

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).

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).

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  print("Device: {}".format(args['device']))

  # Load the data
  data = torch.load("acm.pkl")

  # Message types
  message_type_1 = ("paper", "author", "paper")
  message_type_2 = ("paper", "subject", "paper")

  # Dictionary of edge indices
  edge_index = {}
  edge_index[message_type_1] = data['pap']
  edge_index[message_type_2] = data['psp']

  # Dictionary of node features
  node_feature = {}
  node_feature["paper"] = data['feature']

  # Dictionary of node labels
  node_label = {}
  node_label["paper"] = data['label']

  # Load the train, validation and test indices
  train_idx = {"paper": data['train_idx'].to(args['device'])}
  val_idx = {"paper": data['val_idx'].to(args['device'])}
  test_idx = {"paper": data['test_idx'].to(args['device'])}

  # Construct a deepsnap tensor backend HeteroGraph
  hetero_graph = HeteroGraph(
      node_feature=node_feature,
      node_label=node_label,
      edge_index=edge_index,
      directed=True
  )

  print(f"ACM heterogeneous graph: {hetero_graph.num_nodes()} nodes, {hetero_graph.num_edges()} edges")

  # Node feature and node label to device
  for key in hetero_graph.node_feature:
      hetero_graph.node_feature[key] = hetero_graph.node_feature[key].to(args['device'])
  for key in hetero_graph.node_label:
      hetero_graph.node_label[key] = hetero_graph.node_label[key].to(args['device'])

  # Edge_index to sparse tensor and to device
  for key in hetero_graph.edge_index:
      edge_index = hetero_graph.edge_index[key]
      adj = SparseTensor(row=edge_index[0], col=edge_index[1], sparse_sizes=(hetero_graph.num_nodes('paper'), hetero_graph.num_nodes('paper')))
      hetero_graph.edge_index[key] = adj.t().to(args['device'])
  print(hetero_graph.edge_index[message_type_1])
  print(hetero_graph.edge_index[message_type_2])
```

    Device: cuda
    ACM heterogeneous graph: {'paper': 3025} nodes, {('paper', 'author', 'paper'): 26256, ('paper', 'subject', 'paper'): 2207736} edges
    SparseTensor(row=tensor([   0,    0,    0,  ..., 3024, 3024, 3024], device='cuda:0'),
                 col=tensor([   8,   20,   51,  ..., 2948, 2983, 2991], device='cuda:0'),
                 size=(3025, 3025), nnz=26256, density=0.29%)
    SparseTensor(row=tensor([   0,    0,    0,  ..., 3024, 3024, 3024], device='cuda:0'),
                 col=tensor([  75,  434,  534,  ..., 3020, 3021, 3022], device='cuda:0'),
                 size=(3025, 3025), nnz=2207736, density=24.13%)

## Start Training!

Now lets start training!

## Training the Mean Aggregation

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  best_model = None
  best_val = 0

  model = HeteroGNN(hetero_graph, args, aggr="mean").to(args['device'])
  optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])

  for epoch in range(args['epochs']):
      loss = train(model, optimizer, hetero_graph, train_idx)
      accs, best_model, best_val = test(model, hetero_graph, [train_idx, val_idx, test_idx], best_model, best_val)
      print(
          f"Epoch {epoch + 1}: loss {round(loss, 5)}, "
          f"train micro {round(accs[0][0] * 100, 2)}%, train macro {round(accs[0][1] * 100, 2)}%, "
          f"valid micro {round(accs[1][0] * 100, 2)}%, valid macro {round(accs[1][1] * 100, 2)}%, "
          f"test micro {round(accs[2][0] * 100, 2)}%, test macro {round(accs[2][1] * 100, 2)}%"
      )
  best_accs, _, _ = test(best_model, hetero_graph, [train_idx, val_idx, test_idx], save_preds=True, agg_type="Mean")
  print(
      f"Best model: "
      f"train micro {round(best_accs[0][0] * 100, 2)}%, train macro {round(best_accs[0][1] * 100, 2)}%, "
      f"valid micro {round(best_accs[1][0] * 100, 2)}%, valid macro {round(best_accs[1][1] * 100, 2)}%, "
      f"test micro {round(best_accs[2][0] * 100, 2)}%, test macro {round(best_accs[2][1] * 100, 2)}%"
  )
```

    Epoch 1: loss 1.09882, train micro 33.33%, train macro 16.67%, valid micro 33.33%, valid macro 16.67%, test micro 35.81%, test macro 17.58%
    Epoch 2: loss 1.09196, train micro 49.83%, train macro 41.09%, valid micro 38.0%, valid macro 25.46%, test micro 39.95%, test macro 25.92%
    Epoch 3: loss 1.06808, train micro 65.67%, train macro 54.51%, valid micro 65.0%, valid macro 54.34%, test micro 64.47%, test macro 53.97%
    Epoch 4: loss 1.01791, train micro 66.33%, train macro 54.96%, valid micro 66.33%, valid macro 55.31%, test micro 65.04%, test macro 54.28%
    Epoch 5: loss 0.92926, train micro 66.83%, train macro 55.98%, valid micro 66.0%, valid macro 55.0%, test micro 65.36%, test macro 54.43%
    Epoch 6: loss 0.7961, train micro 67.33%, train macro 56.78%, valid micro 66.0%, valid macro 55.0%, test micro 65.41%, test macro 54.42%
    Epoch 7: loss 0.63329, train micro 67.67%, train macro 57.41%, valid micro 66.33%, valid macro 55.64%, test micro 65.46%, test macro 54.44%
    Epoch 8: loss 0.47548, train micro 68.83%, train macro 59.85%, valid micro 66.33%, valid macro 55.73%, test micro 65.55%, test macro 54.57%
    Epoch 9: loss 0.3465, train micro 71.0%, train macro 63.93%, valid micro 66.33%, valid macro 55.73%, test micro 65.55%, test macro 54.53%
    Epoch 10: loss 0.24995, train micro 75.0%, train macro 70.55%, valid micro 68.33%, valid macro 59.8%, test micro 65.74%, test macro 54.86%
    Epoch 11: loss 0.1809, train micro 79.0%, train macro 76.31%, valid micro 73.33%, valid macro 68.61%, test micro 66.21%, test macro 55.85%
    Epoch 12: loss 0.13183, train micro 86.83%, train macro 86.01%, valid micro 80.0%, valid macro 78.11%, test micro 67.34%, test macro 58.23%
    Epoch 13: loss 0.09647, train micro 93.0%, train macro 92.81%, valid micro 86.0%, valid macro 85.19%, test micro 70.02%, test macro 63.13%
    Epoch 14: loss 0.07086, train micro 95.33%, train macro 95.26%, valid micro 90.33%, valid macro 90.0%, test micro 72.61%, test macro 67.57%
    Epoch 15: loss 0.05302, train micro 96.83%, train macro 96.81%, valid micro 91.33%, valid macro 91.09%, test micro 74.73%, test macro 70.96%
    Epoch 16: loss 0.04085, train micro 97.5%, train macro 97.49%, valid micro 93.33%, valid macro 93.18%, test micro 77.13%, test macro 74.36%
    Epoch 17: loss 0.03248, train micro 98.17%, train macro 98.16%, valid micro 93.67%, valid macro 93.52%, test micro 79.01%, test macro 76.93%
    Epoch 18: loss 0.02659, train micro 99.0%, train macro 99.0%, valid micro 94.67%, valid macro 94.59%, test micro 80.24%, test macro 78.56%
    Epoch 19: loss 0.02235, train micro 99.17%, train macro 99.17%, valid micro 95.67%, valid macro 95.63%, test micro 81.22%, test macro 79.85%
    Epoch 20: loss 0.01919, train micro 99.33%, train macro 99.33%, valid micro 96.67%, valid macro 96.65%, test micro 82.12%, test macro 80.95%
    Epoch 21: loss 0.01679, train micro 99.33%, train macro 99.33%, valid micro 97.0%, valid macro 97.0%, test micro 83.06%, test macro 82.1%
    Epoch 22: loss 0.01491, train micro 99.67%, train macro 99.67%, valid micro 97.0%, valid macro 97.0%, test micro 83.53%, test macro 82.69%
    Epoch 23: loss 0.01338, train micro 99.83%, train macro 99.83%, valid micro 97.0%, valid macro 97.0%, test micro 84.19%, test macro 83.49%
    Epoch 24: loss 0.01209, train micro 99.83%, train macro 99.83%, valid micro 97.33%, valid macro 97.33%, test micro 84.61%, test macro 83.98%
    Epoch 25: loss 0.01097, train micro 99.83%, train macro 99.83%, valid micro 97.67%, valid macro 97.66%, test micro 84.85%, test macro 84.27%
    Epoch 26: loss 0.01001, train micro 99.83%, train macro 99.83%, valid micro 97.67%, valid macro 97.66%, test micro 85.13%, test macro 84.58%
    Epoch 27: loss 0.00918, train micro 99.83%, train macro 99.83%, valid micro 98.0%, valid macro 97.99%, test micro 85.41%, test macro 84.91%
    Epoch 28: loss 0.00837, train micro 99.83%, train macro 99.83%, valid micro 98.33%, valid macro 98.33%, test micro 85.32%, test macro 84.81%
    Epoch 29: loss 0.0074, train micro 99.83%, train macro 99.83%, valid micro 98.33%, valid macro 98.33%, test micro 85.46%, test macro 84.98%
    Epoch 30: loss 0.00651, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.33%, test micro 85.46%, test macro 85.02%
    Epoch 31: loss 0.00572, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.51%, test macro 85.09%
    Epoch 32: loss 0.00505, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.69%, test macro 85.3%
    Epoch 33: loss 0.0045, train micro 100.0%, train macro 100.0%, valid micro 98.33%, valid macro 98.34%, test micro 85.6%, test macro 85.2%
    Epoch 34: loss 0.00396, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.55%, test macro 85.21%
    Epoch 35: loss 0.00323, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.65%, test macro 85.32%
    Epoch 36: loss 0.00255, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.69%, test macro 85.37%
    Epoch 37: loss 0.00201, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.51%, test macro 85.19%
    Epoch 38: loss 0.00161, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.51%, test macro 85.21%
    Epoch 39: loss 0.00134, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.41%, test macro 85.13%
    Epoch 40: loss 0.00115, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.27%, test macro 85.0%
    Epoch 41: loss 0.00102, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.18%, test macro 84.93%
    Epoch 42: loss 0.00093, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.18%, test macro 84.94%
    Epoch 43: loss 0.00086, train micro 100.0%, train macro 100.0%, valid micro 98.0%, valid macro 98.0%, test micro 85.22%, test macro 85.02%
    Epoch 44: loss 0.0008, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.67%, test micro 85.18%, test macro 84.99%
    Epoch 45: loss 0.00076, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.67%, test micro 85.22%, test macro 85.04%
    Epoch 46: loss 0.00072, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.18%, test macro 85.01%
    Epoch 47: loss 0.00069, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.22%, test macro 85.06%
    Epoch 48: loss 0.00066, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.66%, test micro 85.27%, test macro 85.1%
    Epoch 49: loss 0.00063, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.32%, test macro 85.16%
    Epoch 50: loss 0.00061, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.36%, test macro 85.21%
    Epoch 51: loss 0.00059, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.41%, test macro 85.27%
    Epoch 52: loss 0.00057, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.46%, test macro 85.32%
    Epoch 53: loss 0.00055, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 85.51%, test macro 85.37%
    Epoch 54: loss 0.00054, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 85.46%, test macro 85.32%
    Epoch 55: loss 0.00052, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 85.51%, test macro 85.37%
    Epoch 56: loss 0.00051, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.66%, test micro 85.36%, test macro 85.22%
    Epoch 57: loss 0.0005, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.55%, test macro 85.43%
    Epoch 58: loss 0.00048, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.55%, test macro 85.43%
    Epoch 59: loss 0.00047, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.49%
    Epoch 60: loss 0.00046, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.51%
    Epoch 61: loss 0.00045, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.69%, test macro 85.61%
    Epoch 62: loss 0.00044, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.69%, test macro 85.61%
    Epoch 63: loss 0.00043, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.69%, test macro 85.61%
    Epoch 64: loss 0.00042, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.69%, test macro 85.62%
    Epoch 65: loss 0.00042, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.74%, test macro 85.67%
    Epoch 66: loss 0.00041, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.58%
    Epoch 67: loss 0.0004, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.58%
    Epoch 68: loss 0.0004, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.69%, test macro 85.64%
    Epoch 69: loss 0.00039, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.79%, test macro 85.74%
    Epoch 70: loss 0.00038, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.84%, test macro 85.79%
    Epoch 71: loss 0.00038, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.84%, test macro 85.79%
    Epoch 72: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.84%, test macro 85.79%
    Epoch 73: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.74%, test macro 85.7%
    Epoch 74: loss 0.00036, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.61%
    Epoch 75: loss 0.00036, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.61%
    Epoch 76: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.74%, test macro 85.71%
    Epoch 77: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.61%
    Epoch 78: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.74%, test macro 85.71%
    Epoch 79: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.74%, test macro 85.71%
    Epoch 80: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.74%, test macro 85.71%
    Epoch 81: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.79%, test macro 85.76%
    Epoch 82: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.69%, test macro 85.67%
    Epoch 83: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.63%
    Epoch 84: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.63%
    Epoch 85: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.63%
    Epoch 86: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.65%, test macro 85.63%
    Epoch 87: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.58%
    Epoch 88: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.58%
    Epoch 89: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.58%
    Epoch 90: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.58%
    Epoch 91: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.58%
    Epoch 92: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.6%, test macro 85.58%
    Epoch 93: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.99%, test micro 85.55%, test macro 85.54%
    Epoch 94: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.55%, test macro 85.54%
    Epoch 95: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.51%, test macro 85.5%
    Epoch 96: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.55%, test macro 85.55%
    Epoch 97: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.55%, test macro 85.55%
    Epoch 98: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.6%, test macro 85.59%
    Epoch 99: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.65%, test macro 85.64%
    Epoch 100: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.6%, test macro 85.59%
    Saving Heterogeneous Node Prediction Model Predictions with Agg: Mean

    Best model: train micro 99.83%, train macro 99.83%, valid micro 98.33%, valid macro 98.33%, test micro 85.32%, test macro 84.81%

## Question 2.1: What is your maximum test set **micro** F1 score for the best_model when using mean aggregation? (10 points)

## Question 2.2: What is your maximum test set **macro** F1 score for the best_model when using the mean aggregation? (10 points)

## Training the Attention Aggregation

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  best_model = None
  best_val = 0

  output_size = hetero_graph.num_node_labels('paper')
  model = HeteroGNN(hetero_graph, args, aggr="attn").to(args['device'])
  optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])

  for epoch in range(args['epochs']):
      loss = train(model, optimizer, hetero_graph, train_idx)
      accs, best_model, best_val = test(model, hetero_graph, [train_idx, val_idx, test_idx], best_model, best_val)
      print(
          f"Epoch {epoch + 1}: loss {round(loss, 5)}, "
          f"train micro {round(accs[0][0] * 100, 2)}%, train macro {round(accs[0][1] * 100, 2)}%, "
          f"valid micro {round(accs[1][0] * 100, 2)}%, valid macro {round(accs[1][1] * 100, 2)}%, "
          f"test micro {round(accs[2][0] * 100, 2)}%, test macro {round(accs[2][1] * 100, 2)}%"
      )
  best_accs, _, _ = test(best_model, hetero_graph, [train_idx, val_idx, test_idx], save_preds=True, agg_type="Attention")
  print(
      f"Best model: "
      f"train micro {round(best_accs[0][0] * 100, 2)}%, train macro {round(best_accs[0][1] * 100, 2)}%, "
      f"valid micro {round(best_accs[1][0] * 100, 2)}%, valid macro {round(best_accs[1][1] * 100, 2)}%, "
      f"test micro {round(best_accs[2][0] * 100, 2)}%, test macro {round(best_accs[2][1] * 100, 2)}%"
  )
```

    Epoch 1: loss 1.0995, train micro 33.33%, train macro 16.67%, valid micro 33.33%, valid macro 16.67%, test micro 32.89%, test macro 16.5%
    Epoch 2: loss 1.08985, train micro 84.5%, train macro 84.91%, valid micro 69.67%, valid macro 69.76%, test micro 61.04%, test macro 59.79%
    Epoch 3: loss 1.05637, train micro 75.17%, train macro 71.33%, valid micro 74.33%, valid macro 70.14%, test micro 65.22%, test macro 57.87%
    Epoch 4: loss 0.99032, train micro 67.67%, train macro 58.05%, valid micro 66.33%, valid macro 55.31%, test micro 65.08%, test macro 54.36%
    Epoch 5: loss 0.87888, train micro 66.83%, train macro 56.21%, valid micro 66.33%, valid macro 55.31%, test micro 65.18%, test macro 54.34%
    Epoch 6: loss 0.72638, train micro 66.33%, train macro 55.12%, valid micro 66.33%, valid macro 55.31%, test micro 65.22%, test macro 54.37%
    Epoch 7: loss 0.56652, train micro 66.33%, train macro 55.12%, valid micro 66.33%, valid macro 55.31%, test micro 65.22%, test macro 54.38%
    Epoch 8: loss 0.43753, train micro 66.33%, train macro 55.08%, valid micro 66.33%, valid macro 55.31%, test micro 65.22%, test macro 54.38%
    Epoch 9: loss 0.34556, train micro 66.67%, train macro 55.69%, valid micro 66.0%, valid macro 55.0%, test micro 65.22%, test macro 54.34%
    Epoch 10: loss 0.27648, train micro 67.5%, train macro 57.2%, valid micro 66.33%, valid macro 55.73%, test micro 65.32%, test macro 54.42%
    Epoch 11: loss 0.21962, train micro 69.83%, train macro 61.82%, valid micro 67.33%, valid macro 57.77%, test micro 65.74%, test macro 54.98%
    Epoch 12: loss 0.17146, train micro 76.5%, train macro 72.86%, valid micro 76.0%, valid macro 72.76%, test micro 66.96%, test macro 57.56%
    Epoch 13: loss 0.13111, train micro 84.67%, train macro 83.52%, valid micro 84.33%, valid macro 83.44%, test micro 69.79%, test macro 62.66%
    Epoch 14: loss 0.09864, train micro 93.33%, train macro 93.14%, valid micro 88.67%, valid macro 88.16%, test micro 73.27%, test macro 68.31%
    Epoch 15: loss 0.07401, train micro 95.17%, train macro 95.08%, valid micro 91.33%, valid macro 91.02%, test micro 75.44%, test macro 71.72%
    Epoch 16: loss 0.0561, train micro 96.5%, train macro 96.46%, valid micro 92.67%, valid macro 92.49%, test micro 76.75%, test macro 73.8%
    Epoch 17: loss 0.043, train micro 97.33%, train macro 97.31%, valid micro 93.0%, valid macro 92.96%, test micro 77.93%, test macro 75.59%
    Epoch 18: loss 0.03316, train micro 98.17%, train macro 98.16%, valid micro 94.0%, valid macro 93.97%, test micro 78.54%, test macro 76.77%
    Epoch 19: loss 0.02601, train micro 99.33%, train macro 99.33%, valid micro 95.33%, valid macro 95.34%, test micro 79.11%, test macro 77.89%
    Epoch 20: loss 0.02102, train micro 99.5%, train macro 99.5%, valid micro 95.0%, valid macro 95.02%, test micro 79.81%, test macro 78.92%
    Epoch 21: loss 0.01751, train micro 99.83%, train macro 99.83%, valid micro 95.0%, valid macro 95.02%, test micro 79.81%, test macro 79.07%
    Epoch 22: loss 0.01492, train micro 100.0%, train macro 100.0%, valid micro 94.33%, valid macro 94.36%, test micro 80.28%, test macro 79.72%
    Epoch 23: loss 0.0129, train micro 100.0%, train macro 100.0%, valid micro 94.0%, valid macro 93.98%, test micro 80.89%, test macro 80.5%
    Epoch 24: loss 0.01127, train micro 100.0%, train macro 100.0%, valid micro 94.67%, valid macro 94.62%, test micro 81.6%, test macro 81.3%
    Epoch 25: loss 0.00981, train micro 100.0%, train macro 100.0%, valid micro 95.67%, valid macro 95.63%, test micro 82.54%, test macro 82.3%
    Epoch 26: loss 0.00848, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 95.95%, test micro 83.48%, test macro 83.26%
    Epoch 27: loss 0.00737, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.29%, test micro 84.09%, test macro 83.89%
    Epoch 28: loss 0.00645, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.29%, test micro 84.52%, test macro 84.32%
    Epoch 29: loss 0.00556, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 96.98%, test micro 85.04%, test macro 84.85%
    Epoch 30: loss 0.00473, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.32%, test micro 85.22%, test macro 85.05%
    Epoch 31: loss 0.00371, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.32%, test micro 85.46%, test macro 85.32%
    Epoch 32: loss 0.00279, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 85.41%, test macro 85.28%
    Epoch 33: loss 0.00216, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 85.79%, test macro 85.68%
    Epoch 34: loss 0.0017, train micro 100.0%, train macro 100.0%, valid micro 97.33%, valid macro 97.33%, test micro 85.46%, test macro 85.34%
    Epoch 35: loss 0.00138, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.67%, test micro 85.18%, test macro 85.06%
    Epoch 36: loss 0.00117, train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.67%, test micro 84.89%, test macro 84.78%
    Epoch 37: loss 0.00102, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 84.8%, test macro 84.69%
    Epoch 38: loss 0.00091, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 84.52%, test macro 84.4%
    Epoch 39: loss 0.00082, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 84.33%, test macro 84.2%
    Epoch 40: loss 0.00076, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 84.09%, test macro 83.97%
    Epoch 41: loss 0.0007, train micro 100.0%, train macro 100.0%, valid micro 97.0%, valid macro 97.0%, test micro 83.81%, test macro 83.68%
    Epoch 42: loss 0.00066, train micro 100.0%, train macro 100.0%, valid micro 96.67%, valid macro 96.67%, test micro 83.72%, test macro 83.58%
    Epoch 43: loss 0.00062, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.67%, test macro 83.54%
    Epoch 44: loss 0.00059, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.62%, test macro 83.49%
    Epoch 45: loss 0.00056, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.53%, test macro 83.4%
    Epoch 46: loss 0.00053, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.3%
    Epoch 47: loss 0.00051, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.39%, test macro 83.26%
    Epoch 48: loss 0.00049, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.44%, test macro 83.31%
    Epoch 49: loss 0.00047, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.36%, test micro 83.53%, test macro 83.41%
    Epoch 50: loss 0.00046, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.36%, test micro 83.48%, test macro 83.35%
    Epoch 51: loss 0.00044, train micro 100.0%, train macro 100.0%, valid micro 96.0%, valid macro 96.01%, test micro 83.53%, test macro 83.4%
    Epoch 52: loss 0.00043, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.34%, test micro 83.62%, test macro 83.5%
    Epoch 53: loss 0.00042, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.62%, test macro 83.5%
    Epoch 54: loss 0.0004, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.62%, test macro 83.5%
    Epoch 55: loss 0.00039, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.53%, test macro 83.41%
    Epoch 56: loss 0.00038, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.48%, test macro 83.37%
    Epoch 57: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.39%, test macro 83.28%
    Epoch 58: loss 0.00037, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.44%, test macro 83.34%
    Epoch 59: loss 0.00036, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.34%, test macro 83.26%
    Epoch 60: loss 0.00035, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.29%, test macro 83.22%
    Epoch 61: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.25%, test macro 83.18%
    Epoch 62: loss 0.00034, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.2%, test macro 83.14%
    Epoch 63: loss 0.00033, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.34%, test macro 83.28%
    Epoch 64: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.25%, test macro 83.19%
    Epoch 65: loss 0.00032, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.2%, test macro 83.15%
    Epoch 66: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.2%, test macro 83.16%
    Epoch 67: loss 0.00031, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 83.11%, test macro 83.07%
    Epoch 68: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.96%, test macro 82.95%
    Epoch 69: loss 0.0003, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.92%, test macro 82.9%
    Epoch 70: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.92%, test macro 82.9%
    Epoch 71: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.82%, test macro 82.8%
    Epoch 72: loss 0.00029, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.73%, test macro 82.72%
    Epoch 73: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.82%, test macro 82.82%
    Epoch 74: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.78%, test macro 82.77%
    Epoch 75: loss 0.00028, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.78%, test macro 82.77%
    Epoch 76: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.73%, test macro 82.73%
    Epoch 77: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.78%, test macro 82.78%
    Epoch 78: loss 0.00027, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.78%, test macro 82.78%
    Epoch 79: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.73%, test macro 82.73%
    Epoch 80: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.68%, test macro 82.69%
    Epoch 81: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.68%, test macro 82.69%
    Epoch 82: loss 0.00026, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.68%, test macro 82.69%
    Epoch 83: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.78%, test macro 82.79%
    Epoch 84: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.73%, test macro 82.75%
    Epoch 85: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.73%, test macro 82.75%
    Epoch 86: loss 0.00025, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.68%, test macro 82.7%
    Epoch 87: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.68%, test macro 82.7%
    Epoch 88: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.68%, test macro 82.7%
    Epoch 89: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.64%, test macro 82.66%
    Epoch 90: loss 0.00024, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.54%, test macro 82.57%
    Epoch 91: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.54%, test macro 82.57%
    Epoch 92: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.54%, test macro 82.57%
    Epoch 93: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.59%, test macro 82.62%
    Epoch 94: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.59%, test macro 82.62%
    Epoch 95: loss 0.00023, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.59%, test macro 82.62%
    Epoch 96: loss 0.00022, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.59%, test macro 82.62%
    Epoch 97: loss 0.00022, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.59%, test macro 82.62%
    Epoch 98: loss 0.00022, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.59%, test macro 82.62%
    Epoch 99: loss 0.00022, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.64%, test macro 82.67%
    Epoch 100: loss 0.00022, train micro 100.0%, train macro 100.0%, valid micro 96.33%, valid macro 96.33%, test micro 82.64%, test macro 82.67%
    Saving Heterogeneous Node Prediction Model Predictions with Agg: Attention

    Best model: train micro 100.0%, train macro 100.0%, valid micro 97.67%, valid macro 97.67%, test micro 85.18%, test macro 85.06%

## Question 2.3: What is your maximum test set **micro** F1 score for the best_model when using the attention aggregation? (4 points)

## Question 2.4: What is your maximum test set **macro** F1 score for the best_model when using the attention aggregation? (4 points)

## Attention for each Message Type

Through message type level attention we can learn which message type is more important to which layer.

Here we will print out and show that each layer pay how much attention on each message type.

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  if model.convs1.alpha is not None and model.convs2.alpha is not None:
      for idx, message_type in model.convs1.mapping.items():
          print(f"Layer 1 has attention {model.convs1.alpha[idx]} on message type {message_type}")
      for idx, message_type in model.convs2.mapping.items():
          print(f"Layer 2 has attention {model.convs2.alpha[idx]} on message type {message_type}")
```

    Layer 1 has attention tensor([0.4995, 0.4823, 0.4894,  ..., 0.4992, 0.4317, 0.4156], device='cuda:0') on message type ('paper', 'author', 'paper')
    Layer 1 has attention tensor([0.5005, 0.5177, 0.5106,  ..., 0.5008, 0.5683, 0.5844], device='cuda:0') on message type ('paper', 'subject', 'paper')
    Layer 2 has attention tensor([0.4713, 0.4562, 0.4442,  ..., 0.4948, 0.4706, 0.4796], device='cuda:0') on message type ('paper', 'author', 'paper')
    Layer 2 has attention tensor([0.5287, 0.5438, 0.5558,  ..., 0.5052, 0.5294, 0.5204], device='cuda:0') on message type ('paper', 'subject', 'paper')

# Submission

You will need to submit three files on Gradescope to complete this notebook.

1.   Your completed *XCS224W_Colab4.ipynb*. From the "File" menu select "Download .ipynb" to save a local copy of your completed Colab.
2.  *ACM-Node-MeanAgg.csv*
3.  *ACM-Node-AttentionAgg.csv*

Download the csv files by selecting the *Folder* icon on the left panel.

To submit your work, zip the files downloaded in steps 1-3 above and submit to gradescope. **NOTE:** DO NOT rename any of the downloaded files.


```python
import numpy as np
```

```python
t = 24
alpha = 0.05
1.7 * np.sqrt((np.log(np.log(2 * t)) + 0.72 * np.log(5.2 / alpha)) / t)
```

    0.752104016744892

```python
t = 14
1.7 * np.sqrt(t * (np.log10(np.log10(2 * t))))
```

    2.5484261237266748

```python

```

```python
# Add this to python2 code to make life easier -- safe to use in python3 as well.
from __future__ import absolute_import, division, print_function
```

```python
import numpy as np
# don't do:
# from numpy import *
```

```python
max('a')
```

```python
np.max('a')
```

```python
%matplotlib inline
%config InlineBackend.figure_format='retina'
```

```python
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
from pivottablejs import pivot_ui
import sys
import warnings
warnings.filterwarnings('ignore')
```

```python
first = 1
```

```python
second = 2
```

```python
third = 3
```

```python
# df = pd.read_csv("../data/mps.csv", encoding="ISO-8859-1")
# canadian_politics = pd.read_csv("../data/montreal_2014.csv",)
canadian_politics = pd.read_csv("../data/mps2.csv",)
```

```python
# recommend using .head()
canadian_politics.head(10)
```

```python
sns.distplot(canadian_politics['Age'].dropna())
```

```python
sns.distplot(canadian_politics['Age'])
```

```python
sns.distplot(canadian_politics['Age'].dropna())
```

```python
sns.set_context('poster', font_scale=1.3)
```

```python
fig, ax = plt.subplots(figsize=(12, 8))
sns.distplot(canadian_politics['Age'].dropna())
fig.tight_layout()
```

```python
sns.distplot(canadian_politics['Age'].dropna());
```

```python
pivot_ui(canadian_politics)
```

```python
import numpy as np
```

```python
np.linspace?
```

```python
np.linspace??
```

```python
np.linspace??
```

```python
!subl ~/miniconda3/envs/dspy3/lib/python3.6/site-packages/numpy/core/function_base.py
```

```python
import textwrap
def example_function():
    """Docstring for example function"""

    print(textwrap.dedent("""
    This is a multi-lined string
    that I want to write inside of a function.
    Notice what happens when I print this.
        And when something is indented more."""))

example_function()
```

```python
# But first find and replace
def silly_function(xval):
    """Takes a value and returns the value."""
    xval_sq = xval ** 2.0
    1 + 10
    xval_abs = np.sqrt(xval_sq)
    return xval_abs
```

```python
silly_function(2,)
```

```python
silly_function(-2,)
```

```python
silly_function?
```

```python
silly_function??
```

```python
# But first find and replace
def silly_function(xval):
    """Takes a value and returns the value."""
    xval_sq = xval ** 2.0
    1 + 12
    xval_abs = np.sqrt(xval_sq)
    return xval_abs
```

```python
silly_function??
```

```python
!ls ../data/
```

```python
coal_years = !ls ../data/coal_prod_20*.csv
```

```python
coal_years
```

```python
example_dict = {}

# Indent/dedent/comment
for index in range(5):
    example_dict["float_one"] = 1.0
    example_dict["float_two"] = 2.0
    example_dict["float_three"] = 3.0
    example_dict["float_four"] = 4.0
```

```python
example["one_better_neat"] = 1
example["two_better_neat"] = 2
example["three_better_neat"] = 3
example["four_better_neat"] = 4
```

```python
example_ = {}

# Indent/dedent/comment
for index in range(5):
    example_["float_one"] = 1.0
    example_["float_two"] = 2.0
    example_["float_three"] = 3.0
    example_["float_four"] = 4.0
```

```python
example["one_better_neat"] = 1
example["two_better_neat"] = 2
example["three_better_neat"] = 3
example["four_better_neat"] = 4
```

```python
example = {}

# Indent/dedent/comment
for index in range(5):
    example["float_one"] = 1.0
    example["float_two"] = 2.0
    example["float_three"] = 3.0
    example["float_four"] = 4.0
```

```python
example["one_better_neat"] = 1
example["two_better_neat"] = 2
example["three_better_neat"] = 3
example["four_better_neat"] = 4
```

```latex
%%latex

If you want to get crazier$\ldots$

\begin{equation}
\oint_S {E_n dA = \frac{1}{{\varepsilon _0 }}} Q_\textrm{inside}
\end{equation}
```

```python
%%python2
print "hi"
```

```bash
%%bash
wget http://www.ast.cam.ac.uk/~rfc/vpfit11.1.tar.gz
mkdir -p vpfit11
cd vpfit11
tar -xvzf ../vpfit11.1.tar.gz
```

```python
normal_argument = 12.4
second_argument = 98.4

arg_with_spaces = "the secret to life"
```

```bash
%%bash -s {normal_argument} {second_argument}
echo "This script knows the value of the argument: $1"
echo "It also has no trouble with the second argument: $2"
```

```bash
%%bash -s "$arg_with_spaces"
echo "This bash script knows $1."
```

```python
ls
```

```python
!ls
```

```python
%ls
```

```python
ls vpfit10/
```

```python
tailthing = "*.ipynb"
```

```python
tailthing
```

```python
!ls {tailthing}
```

```python
output = !ls
```

```python
output
```

```python
!pwd
```

```python
a = 3
a
```

```python
_58
```

```python
_i58
```