CS224W - Colab 2¶

In Colab 2, we will work to construct our own graph neural network using PyTorch Geometric (PyG) and then apply that model on two Open Graph Benchmark (OGB) datasets. These two datasets will be used to benchmark your model’s performance on two different graph-based tasks: 1) node property prediction, predicting properties of single nodes and 2) graph property prediction, predicting properties of entire graphs or subgraphs.

First, we will learn how PyTorch Geometric stores graphs as PyTorch tensors.

Then, we will load and inspect one of the Open Graph Benchmark (OGB) datasets by using the ogb package. OGB is a collection of realistic, large-scale, and diverse benchmark datasets for machine learning on graphs. The ogb package not only provides data loaders for each dataset but also model evaluators.

Lastly, we will build our own graph neural network using PyTorch Geometric. We will then train and evaluate our model on the OGB node property prediction and graph property prediction tasks.

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

We recommend you save a copy of this colab in your drive so you don’t lose progress!

Have fun and good luck on Colab 2 :)

Device¶

You might need to use a GPU for this Colab to run quickly.

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

Setup¶

As discussed in Colab 0, the installation of PyG on Colab can be a little bit tricky. First let us check which version of PyTorch you are running

[1]:

import torch
import os
print("PyTorch has version {}".format(torch.__version__))

PyTorch has version 1.11.0

Download the necessary packages for PyG. Make sure that your version of torch matches the output from the cell above. In case of any issues, more information can be found on the PyG’s installation page.

[2]:

# Install torch geometric
# 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 ogb

1) PyTorch Geometric (Datasets and Data)¶

PyTorch Geometric has two classes for storing and/or transforming graphs into tensor format. One is torch_geometric.datasets, which contains a variety of common graph datasets. Another is torch_geometric.data, which provides the data handling of graphs in PyTorch tensors.

In this section, we will learn how to use torch_geometric.datasets and torch_geometric.data together.

PyG Datasets¶

The torch_geometric.datasets class has many common graph datasets. Here we will explore its usage through one example dataset.

[3]:

from torch_geometric.datasets import TUDataset

if 'IS_GRADESCOPE_ENV' not in os.environ:
  root = './enzymes'
  name = 'ENZYMES'

  # The ENZYMES dataset
  pyg_dataset= TUDataset(root, name)

  # You will find that there are 600 graphs in this dataset
  print(pyg_dataset)

ENZYMES(600)

Question 1: What is the number of classes and number of features in the ENZYMES dataset? (5 points)¶

[4]:

def get_num_classes(pyg_dataset):
  # TODO: Implement a function that takes a PyG dataset object
  # and returns the number of classes for that dataset.

  num_classes = 0

  ############# Your code here ############
  ## (~1 line of code)
  ## Note
  ## 1. Colab autocomplete functionality might be useful.

  num_classes = pyg_dataset.num_classes

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

  return num_classes

def get_num_features(pyg_dataset):
  # TODO: Implement a function that takes a PyG dataset object
  # and returns the number of features for that dataset.

  num_features = 0

  ############# Your code here ############
  ## (~1 line of code)
  ## Note
  ## 1. Colab autocomplete functionality might be useful.

  num_features = pyg_dataset.num_features

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

  return num_features

if 'IS_GRADESCOPE_ENV' not in os.environ:
  num_classes = get_num_classes(pyg_dataset)
  num_features = get_num_features(pyg_dataset)
  print("{} dataset has {} classes".format(name, num_classes))
  print("{} dataset has {} features".format(name, num_features))

ENZYMES dataset has 6 classes
ENZYMES dataset has 3 features

PyG Data¶

Each PyG dataset stores a list of torch_geometric.data.Data objects, where each torch_geometric.data.Data object represents a graph. We can easily get the Data object by indexing into the dataset.

For more information such as what is stored in the Data object, please refer to the documentation.

Question 2: What is the label of the graph with index 100 in the ENZYMES dataset? (5 points)¶

[5]:

def get_graph_class(pyg_dataset, idx):
  # TODO: Implement a function that takes a PyG dataset object,
  # an index of a graph within the dataset, and returns the class/label
  # of the graph (as an integer).

  label = -1

  ############# Your code here ############
  ## (~1 line of code)

  label = pyg_dataset[idx].y

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

  return label

# Here pyg_dataset is a dataset for graph classification
if 'IS_GRADESCOPE_ENV' not in os.environ:
  graph_0 = pyg_dataset[0]
  print(graph_0)
  idx = 100
  label = get_graph_class(pyg_dataset, idx)
  print('Graph with index {} has label {}'.format(idx, label))

Data(edge_index=[2, 168], x=[37, 3], y=[1])
Graph with index 100 has label tensor([4])

Question 3: How many edges does the graph with index 200 have? (5 points)¶

[6]:

def get_graph_num_edges(pyg_dataset, idx):
  # TODO: Implement a function that takes a PyG dataset object,
  # the index of a graph in the dataset, and returns the number of
  # edges in the graph (as an integer). You should not count an edge
  # twice if the graph is undirected. For example, in an undirected
  # graph G, if two nodes v and u are connected by an edge, this edge
  # should only be counted once.

  num_edges = 0

  ############# Your code here ############
  ## Note:
  ## 1. You can't return the data.num_edges directly
  ## 2. We assume the graph is undirected
  ## 3. Look at the PyG dataset built in functions
  ## (~4 lines of code)

  num_edges = pyg_dataset[idx].edge_index.size()[1] // 2

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

  return num_edges

if 'IS_GRADESCOPE_ENV' not in os.environ:
  idx = 200
  num_edges = get_graph_num_edges(pyg_dataset, idx)
  print('Graph with index {} has {} edges'.format(idx, num_edges))

Graph with index 200 has 53 edges

2) Open Graph Benchmark (OGB)¶

The Open Graph Benchmark (OGB) is a collection of realistic, large-scale, and diverse benchmark datasets for machine learning on graphs. Its datasets are automatically downloaded, processed, and split using the OGB Data Loader. The model performance can then be evaluated by using the OGB Evaluator in a unified manner.

Dataset and Data¶

OGB also supports PyG dataset and data classes. Here we take a look on the ogbn-arxiv dataset.

[7]:

import torch_geometric.transforms as T
from ogb.nodeproppred import PygNodePropPredDataset

if 'IS_GRADESCOPE_ENV' not in os.environ:
  dataset_name = 'ogbn-arxiv'
  # Load the dataset and transform it to sparse tensor
  dataset = PygNodePropPredDataset(name=dataset_name,
                                  transform=T.ToSparseTensor())
  print('The {} dataset has {} graph'.format(dataset_name, len(dataset)))

  # Extract the graph
  data = dataset[0]
  print(data)

WARNING:root:The OGB package is out of date. Your version is 1.3.3, while the latest version is 1.3.4.

The ogbn-arxiv dataset has 1 graph
Data(num_nodes=169343, x=[169343, 128], node_year=[169343, 1], y=[169343, 1], adj_t=[169343, 169343, nnz=1166243])

Question 4: How many features are in the ogbn-arxiv graph? (5 points)¶

[8]:

def graph_num_features(data):
  # TODO: Implement a function that takes a PyG data object,
  # and returns the number of features in the graph (as an integer).

  num_features = 0

  ############# Your code here ############
  ## (~1 line of code)
  num_features = data.num_features
  #########################################

  return num_features

if 'IS_GRADESCOPE_ENV' not in os.environ:
  num_features = graph_num_features(data)
  print('The graph has {} features'.format(num_features))

The graph has 128 features

3) GNN: Node Property Prediction¶

In this section we will build our first graph neural network using PyTorch Geometric. Then we will apply it to the task of node property prediction (node classification).

Specifically, we will use GCN as the foundation for your graph neural network (Kipf et al. (2017)). To do so, we will work with PyG’s built-in GCNConv layer.

Setup¶

[9]:

import torch
import pandas as pd
import torch.nn.functional as F
print(torch.__version__)

# The PyG built-in GCNConv
from torch_geometric.nn import GCNConv

import torch_geometric.transforms as T
from ogb.nodeproppred import PygNodePropPredDataset, Evaluator

1.11.0

Load and Preprocess the Dataset¶

[10]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
  dataset_name = 'ogbn-arxiv'
  dataset = PygNodePropPredDataset(name=dataset_name,
                                  transform=T.ToSparseTensor())
  data = dataset[0]

  # Make the adjacency matrix to symmetric
  data.adj_t = data.adj_t.to_symmetric()

  device = 'cuda' if torch.cuda.is_available() else 'cpu'

  # If you use GPU, the device should be cuda
  print('Device: {}'.format(device))

  data = data.to(device)
  split_idx = dataset.get_idx_split()
  train_idx = split_idx['train'].to(device)

Device: cuda

GCN Model¶

Now we will implement our GCN model!

Please follow the figure below to implement the forward function.

test

[11]:

class GCN(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers,
                 dropout, return_embeds=False):
        # TODO: Implement a function that initializes self.convs,
        # self.bns, and self.softmax.

        super(GCN, self).__init__()

        # A list of GCNConv layers
        self.convs = None

        # A list of 1D batch normalization layers
        self.bns = None

        # The log softmax layer
        self.softmax = None

        ############# Your code here ############
        ## Note:
        ## 1. You should use torch.nn.ModuleList for self.convs and self.bns
        ## 2. self.convs has num_layers GCNConv layers
        ## 3. self.bns has num_layers - 1 BatchNorm1d layers
        ## 4. You should use torch.nn.LogSoftmax for self.softmax
        ## 5. The parameters you can set for GCNConv include 'in_channels' and
        ## 'out_channels'. For more information please refer to the documentation:
        ## https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GCNConv
        ## 6. The only parameter you need to set for BatchNorm1d is 'num_features'
        ## For more information please refer to the documentation:
        ## https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html
        ## (~10 lines of code)

        # 1. use torch.nn.ModuleList for self.convs
        # 2. self.convs has num_layers GCNConv layers
        self.convs = torch.nn.ModuleList()
        for i in range(num_layers - 1):
            self.convs.append(GCNConv(input_dim, hidden_dim))
            input_dim = hidden_dim
        self.convs.append(GCNConv(hidden_dim, output_dim))

        # 1. use torch.nn.ModuleList for self.bns
        # 3. self.bns has num_layers - 1 BatchNorm1d layers
        self.bns = torch.nn.ModuleList([
            torch.nn.BatchNorm1d(num_features=hidden_dim)
                for i in range(num_layers-1)
        ])

        # 4. torch.nn.LogSoftmax for self.softmax
        self.softmax = torch.nn.LogSoftmax()

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

        # Probability of an element getting zeroed
        self.dropout = dropout

        # Skip classification layer and return node embeddings
        self.return_embeds = return_embeds

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        for bn in self.bns:
            bn.reset_parameters()

    def forward(self, x, adj_t):
        # TODO: Implement a function that takes the feature tensor x and
        # edge_index tensor adj_t and returns the output tensor as
        # shown in the figure.

        out = None

        ############# Your code here ############
        ## Note:
        ## 1. Construct the network as shown in the figure
        ## 2. torch.nn.functional.relu and torch.nn.functional.dropout are useful
        ## For more information please refer to the documentation:
        ## https://pytorch.org/docs/stable/nn.functional.html
        ## 3. Don't forget to set F.dropout training to self.training
        ## 4. If return_embeds is True, then skip the last softmax layer
        ## (~7 lines of code)


        for conv, bn in zip(self.convs[:-1], self.bns):
            x = F.relu(bn(conv(x, adj_t)))

            if self.training:

                # 3. F.dropout training
                x = F.dropout(x, p=self.dropout)

        # the last layer
        x = self.convs[-1](x, adj_t)

        # 4. If return_embeds is True, then skip the last softmax layer
        out = x if self.return_embeds else self.softmax(x)

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

        return out

[12]:

def train(model, data, train_idx, optimizer, loss_fn):
    # TODO: Implement a function that trains the model by
    # using the given optimizer and loss_fn.
    model.train()
    loss = 0

    ############# Your code here ############
    ## Note:
    ## 1. Zero grad the optimizer
    ## 2. Feed the data into the model
    ## 3. Slice the model output and label by train_idx
    ## 4. Feed the sliced output and label to loss_fn
    ## (~4 lines of code)

    # 1. Zero grad the optimizer
    optimizer.zero_grad()

    # 2. Feed the data into the model
    output = model(data.x, data.adj_t)

    # 3. Slice the model output and label by train_idx
    output = output[train_idx]
    label = data.y[train_idx,0]

    # 4. Feed the sliced output and label to loss_fn
    loss = loss_fn(output, label)

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

    loss.backward()
    optimizer.step()

    return loss.item()

[13]:

# Test function here
@torch.no_grad()
def test(model, data, split_idx, evaluator, save_model_results=False):
    # TODO: Implement a function that tests the model by
    # using the given split_idx and evaluator.
    model.eval()

    # The output of model on all data
    out = None

    ############# Your code here ############
    ## (~1 line of code)
    ## Note:
    ## 1. No index slicing here
    out = model(data.x, data.adj_t)
    #########################################

    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': data.y[split_idx['train']],
        'y_pred': y_pred[split_idx['train']],
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': data.y[split_idx['valid']],
        'y_pred': y_pred[split_idx['valid']],
    })['acc']
    test_acc = evaluator.eval({
        'y_true': data.y[split_idx['test']],
        'y_pred': y_pred[split_idx['test']],
    })['acc']

    if save_model_results:
      print ("Saving Model Predictions")

      data = {}
      data['y_pred'] = y_pred.view(-1).cpu().detach().numpy()

      df = pd.DataFrame(data=data)
      # Save locally as csv
      df.to_csv('ogbn-arxiv_node.csv', sep=',', index=False)


    return train_acc, valid_acc, test_acc

[14]:

# Please do not change the args
if 'IS_GRADESCOPE_ENV' not in os.environ:
  args = {
      'device': device,
      'num_layers': 3,
      'hidden_dim': 256,
      'dropout': 0.5,
      'lr': 0.01,
      'epochs': 100,
  }
  args

[15]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
  model = GCN(data.num_features, args['hidden_dim'],
              dataset.num_classes, args['num_layers'],
              args['dropout']).to(device)
  evaluator = Evaluator(name='ogbn-arxiv')

[16]:

import copy
if 'IS_GRADESCOPE_ENV' not in os.environ:
  # reset the parameters to initial random value
  model.reset_parameters()

  optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])
  loss_fn = F.nll_loss

  best_model = None
  best_valid_acc = 0

  for epoch in range(1, 1 + args["epochs"]):
    loss = train(model, data, train_idx, optimizer, loss_fn)
    result = test(model, data, split_idx, evaluator)
    train_acc, valid_acc, test_acc = result
    if valid_acc > best_valid_acc:
        best_valid_acc = valid_acc
        best_model = copy.deepcopy(model)
    print(f'Epoch: {epoch:02d}, '
          f'Loss: {loss:.4f}, '
          f'Train: {100 * train_acc:.2f}%, '
          f'Valid: {100 * valid_acc:.2f}% '
          f'Test: {100 * test_acc:.2f}%')

/home/user/anaconda3/envs/gnn/lib/python3.7/site-packages/ipykernel_launcher.py:94: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.

Epoch: 01, Loss: 4.1014, Train: 15.53%, Valid: 24.70% Test: 22.59%
Epoch: 02, Loss: 2.3698, Train: 26.75%, Valid: 24.17% Test: 28.93%
Epoch: 03, Loss: 1.9342, Train: 25.97%, Valid: 20.55% Test: 23.35%
Epoch: 04, Loss: 1.7880, Train: 21.39%, Valid: 10.25% Test: 8.01%
Epoch: 05, Loss: 1.6817, Train: 22.82%, Valid: 11.09% Test: 8.84%
Epoch: 06, Loss: 1.5954, Train: 24.65%, Valid: 12.46% Test: 10.05%
Epoch: 07, Loss: 1.5128, Train: 27.32%, Valid: 14.85% Test: 12.47%
Epoch: 08, Loss: 1.4595, Train: 31.10%, Valid: 20.24% Test: 19.60%
Epoch: 09, Loss: 1.4203, Train: 34.55%, Valid: 28.52% Test: 30.24%
Epoch: 10, Loss: 1.3791, Train: 36.61%, Valid: 33.23% Test: 36.03%
Epoch: 11, Loss: 1.3476, Train: 36.50%, Valid: 33.70% Test: 37.30%
Epoch: 12, Loss: 1.3179, Train: 36.97%, Valid: 34.51% Test: 38.37%
Epoch: 13, Loss: 1.2962, Train: 39.60%, Valid: 37.88% Test: 41.85%
Epoch: 14, Loss: 1.2782, Train: 43.22%, Valid: 42.91% Test: 46.87%
Epoch: 15, Loss: 1.2553, Train: 47.85%, Valid: 49.04% Test: 52.45%
Epoch: 16, Loss: 1.2426, Train: 52.28%, Valid: 53.91% Test: 56.81%
Epoch: 17, Loss: 1.2233, Train: 55.10%, Valid: 56.85% Test: 59.06%
Epoch: 18, Loss: 1.2069, Train: 56.66%, Valid: 58.20% Test: 60.50%
Epoch: 19, Loss: 1.1958, Train: 57.76%, Valid: 59.08% Test: 61.39%
Epoch: 20, Loss: 1.1856, Train: 59.15%, Valid: 60.09% Test: 62.15%
Epoch: 21, Loss: 1.1767, Train: 60.45%, Valid: 61.25% Test: 62.77%
Epoch: 22, Loss: 1.1629, Train: 61.26%, Valid: 61.83% Test: 63.14%
Epoch: 23, Loss: 1.1570, Train: 61.59%, Valid: 61.97% Test: 63.71%
Epoch: 24, Loss: 1.1466, Train: 62.14%, Valid: 62.42% Test: 64.01%
Epoch: 25, Loss: 1.1404, Train: 62.98%, Valid: 63.11% Test: 64.44%
Epoch: 26, Loss: 1.1303, Train: 64.11%, Valid: 64.06% Test: 65.18%
Epoch: 27, Loss: 1.1241, Train: 65.25%, Valid: 65.10% Test: 66.09%
Epoch: 28, Loss: 1.1166, Train: 65.59%, Valid: 65.20% Test: 66.43%
Epoch: 29, Loss: 1.1110, Train: 65.57%, Valid: 65.12% Test: 66.60%
Epoch: 30, Loss: 1.1030, Train: 65.21%, Valid: 64.39% Test: 66.29%
Epoch: 31, Loss: 1.0948, Train: 65.35%, Valid: 64.46% Test: 66.34%
Epoch: 32, Loss: 1.0926, Train: 66.65%, Valid: 66.11% Test: 67.28%
Epoch: 33, Loss: 1.0820, Train: 67.52%, Valid: 67.11% Test: 67.92%
Epoch: 34, Loss: 1.0823, Train: 67.81%, Valid: 67.33% Test: 68.12%
Epoch: 35, Loss: 1.0750, Train: 67.97%, Valid: 67.31% Test: 68.16%
Epoch: 36, Loss: 1.0685, Train: 68.17%, Valid: 67.32% Test: 68.26%
Epoch: 37, Loss: 1.0669, Train: 68.53%, Valid: 67.62% Test: 68.52%
Epoch: 38, Loss: 1.0602, Train: 69.00%, Valid: 68.15% Test: 68.84%
Epoch: 39, Loss: 1.0558, Train: 69.32%, Valid: 68.75% Test: 69.13%
Epoch: 40, Loss: 1.0504, Train: 69.52%, Valid: 68.95% Test: 69.33%
Epoch: 41, Loss: 1.0480, Train: 69.68%, Valid: 69.20% Test: 69.56%
Epoch: 42, Loss: 1.0446, Train: 69.86%, Valid: 69.41% Test: 69.74%
Epoch: 43, Loss: 1.0395, Train: 70.19%, Valid: 69.73% Test: 70.10%
Epoch: 44, Loss: 1.0400, Train: 70.63%, Valid: 70.21% Test: 70.34%
Epoch: 45, Loss: 1.0321, Train: 70.91%, Valid: 70.48% Test: 70.29%
Epoch: 46, Loss: 1.0291, Train: 71.05%, Valid: 70.53% Test: 70.06%
Epoch: 47, Loss: 1.0249, Train: 71.03%, Valid: 70.47% Test: 70.00%
Epoch: 48, Loss: 1.0200, Train: 71.06%, Valid: 70.59% Test: 70.10%
Epoch: 49, Loss: 1.0172, Train: 70.93%, Valid: 70.57% Test: 70.29%
Epoch: 50, Loss: 1.0165, Train: 70.79%, Valid: 70.53% Test: 70.35%
Epoch: 51, Loss: 1.0127, Train: 70.91%, Valid: 70.45% Test: 70.45%
Epoch: 52, Loss: 1.0083, Train: 71.06%, Valid: 70.75% Test: 70.42%
Epoch: 53, Loss: 1.0073, Train: 71.28%, Valid: 70.84% Test: 70.49%
Epoch: 54, Loss: 1.0039, Train: 71.22%, Valid: 70.76% Test: 70.42%
Epoch: 55, Loss: 1.0045, Train: 71.27%, Valid: 70.72% Test: 70.33%
Epoch: 56, Loss: 0.9999, Train: 71.52%, Valid: 70.94% Test: 70.54%
Epoch: 57, Loss: 0.9970, Train: 71.84%, Valid: 71.05% Test: 70.48%
Epoch: 58, Loss: 0.9927, Train: 71.83%, Valid: 70.82% Test: 70.22%
Epoch: 59, Loss: 0.9898, Train: 71.70%, Valid: 70.81% Test: 70.38%
Epoch: 60, Loss: 0.9899, Train: 71.48%, Valid: 70.73% Test: 70.63%
Epoch: 61, Loss: 0.9858, Train: 71.66%, Valid: 70.80% Test: 70.75%
Epoch: 62, Loss: 0.9812, Train: 72.06%, Valid: 71.12% Test: 70.44%
Epoch: 63, Loss: 0.9775, Train: 72.18%, Valid: 71.10% Test: 70.17%
Epoch: 64, Loss: 0.9815, Train: 72.19%, Valid: 71.17% Test: 70.36%
Epoch: 65, Loss: 0.9780, Train: 72.26%, Valid: 71.19% Test: 70.59%
Epoch: 66, Loss: 0.9735, Train: 72.25%, Valid: 71.18% Test: 70.44%
Epoch: 67, Loss: 0.9741, Train: 72.24%, Valid: 71.00% Test: 70.00%
Epoch: 68, Loss: 0.9692, Train: 72.30%, Valid: 71.04% Test: 69.97%
Epoch: 69, Loss: 0.9660, Train: 72.31%, Valid: 71.35% Test: 70.52%
Epoch: 70, Loss: 0.9655, Train: 72.35%, Valid: 71.32% Test: 70.51%
Epoch: 71, Loss: 0.9641, Train: 72.53%, Valid: 71.19% Test: 69.93%
Epoch: 72, Loss: 0.9618, Train: 72.58%, Valid: 71.18% Test: 69.79%
Epoch: 73, Loss: 0.9597, Train: 72.65%, Valid: 71.54% Test: 70.55%
Epoch: 74, Loss: 0.9570, Train: 72.57%, Valid: 71.51% Test: 71.08%
Epoch: 75, Loss: 0.9567, Train: 72.70%, Valid: 71.48% Test: 70.80%
Epoch: 76, Loss: 0.9507, Train: 72.73%, Valid: 70.65% Test: 69.06%
Epoch: 77, Loss: 0.9517, Train: 72.68%, Valid: 70.42% Test: 68.70%
Epoch: 78, Loss: 0.9502, Train: 72.78%, Valid: 71.23% Test: 70.11%
Epoch: 79, Loss: 0.9492, Train: 72.63%, Valid: 71.46% Test: 70.91%
Epoch: 80, Loss: 0.9475, Train: 72.66%, Valid: 71.46% Test: 71.23%
Epoch: 81, Loss: 0.9453, Train: 72.94%, Valid: 71.68% Test: 70.95%
Epoch: 82, Loss: 0.9455, Train: 73.12%, Valid: 71.61% Test: 70.80%
Epoch: 83, Loss: 0.9436, Train: 73.09%, Valid: 71.51% Test: 70.77%
Epoch: 84, Loss: 0.9408, Train: 73.08%, Valid: 71.60% Test: 70.90%
Epoch: 85, Loss: 0.9383, Train: 73.12%, Valid: 71.74% Test: 71.08%
Epoch: 86, Loss: 0.9381, Train: 73.17%, Valid: 71.74% Test: 70.87%
Epoch: 87, Loss: 0.9348, Train: 73.26%, Valid: 71.76% Test: 70.71%
Epoch: 88, Loss: 0.9332, Train: 73.17%, Valid: 71.75% Test: 71.28%
Epoch: 89, Loss: 0.9308, Train: 73.02%, Valid: 71.58% Test: 71.23%
Epoch: 90, Loss: 0.9307, Train: 73.16%, Valid: 71.59% Test: 70.71%
Epoch: 91, Loss: 0.9299, Train: 73.20%, Valid: 71.22% Test: 69.77%
Epoch: 92, Loss: 0.9296, Train: 73.43%, Valid: 71.34% Test: 70.07%
Epoch: 93, Loss: 0.9261, Train: 73.43%, Valid: 71.71% Test: 70.83%
Epoch: 94, Loss: 0.9241, Train: 73.18%, Valid: 71.66% Test: 70.61%
Epoch: 95, Loss: 0.9241, Train: 73.35%, Valid: 71.72% Test: 70.68%
Epoch: 96, Loss: 0.9221, Train: 73.60%, Valid: 71.91% Test: 70.86%
Epoch: 97, Loss: 0.9223, Train: 73.60%, Valid: 71.75% Test: 70.88%
Epoch: 98, Loss: 0.9195, Train: 73.66%, Valid: 71.26% Test: 69.68%
Epoch: 99, Loss: 0.9181, Train: 73.76%, Valid: 71.36% Test: 69.97%
Epoch: 100, Loss: 0.9145, Train: 73.77%, Valid: 71.75% Test: 71.21%

Question 5: What are your `best_model` validation and test accuracies?(20 points)¶

Run the cell below to see the results of your best of model and save your model’s predictions to a file named ogbn-arxiv_node.csv.

You can view this file by clicking on the Folder icon on the left side pannel. As in Colab 1, when you sumbit your assignment, you will have to download this file and attatch it to your submission.

[17]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
  best_result = test(best_model, data, split_idx, evaluator, save_model_results=True)
  train_acc, valid_acc, test_acc = best_result
  print(f'Best model: '
        f'Train: {100 * train_acc:.2f}%, '
        f'Valid: {100 * valid_acc:.2f}% '
        f'Test: {100 * test_acc:.2f}%')

Saving Model Predictions
Best model: Train: 73.60%, Valid: 71.91% Test: 70.86%

/home/user/anaconda3/envs/gnn/lib/python3.7/site-packages/ipykernel_launcher.py:94: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.

4) GNN: Graph Property Prediction¶

In this section we will create a graph neural network for graph property prediction (graph classification).

Load and preprocess the dataset¶

[18]:

from ogb.graphproppred import PygGraphPropPredDataset, Evaluator
from torch_geometric.data import DataLoader
from tqdm.notebook import tqdm

if 'IS_GRADESCOPE_ENV' not in os.environ:
  # Load the dataset
  dataset = PygGraphPropPredDataset(name='ogbg-molhiv')

  device = 'cuda' if torch.cuda.is_available() else 'cpu'
  print('Device: {}'.format(device))

  split_idx = dataset.get_idx_split()

  # Check task type
  print('Task type: {}'.format(dataset.task_type))

Device: cuda
Task type: binary classification

[19]:

# Load the dataset splits into corresponding dataloaders
# We will train the graph classification task on a batch of 32 graphs
# Shuffle the order of graphs for training set
if 'IS_GRADESCOPE_ENV' not in os.environ:
  train_loader = DataLoader(dataset[split_idx["train"]], batch_size=32, shuffle=True, num_workers=0)
  valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=32, shuffle=False, num_workers=0)
  test_loader = DataLoader(dataset[split_idx["test"]], batch_size=32, shuffle=False, num_workers=0)

/home/user/anaconda3/envs/gnn/lib/python3.7/site-packages/torch_geometric/deprecation.py:12: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead
  warnings.warn(out)

[20]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
  # Please do not change the args
  args = {
      'device': device,
      'num_layers': 5,
      'hidden_dim': 256,
      'dropout': 0.5,
      'lr': 0.001,
      'epochs': 30,
  }
  args

Graph Prediction Model¶

Graph Mini-Batching¶

Before diving into the actual model, we introduce the concept of mini-batching with graphs. In order to parallelize the processing of a mini-batch of graphs, PyG combines the graphs into a single disconnected graph data object (torch_geometric.data.Batch). torch_geometric.data.Batch inherits from torch_geometric.data.Data (introduced earlier) and contains an additional attribute called batch.

The batch attribute is a vector mapping each node to the index of its corresponding graph within the mini-batch:

batch = [0, ..., 0, 1, ..., n - 2, n - 1, ..., n - 1]

This attribute is crucial for associating which graph each node belongs to and can be used to e.g. average the node embeddings for each graph individually to compute graph level embeddings.

Implemention¶

Now, we have all of the tools to implement a GCN Graph Prediction model!

We will reuse the existing GCN model to generate node_embeddings and then use Global Pooling over the nodes to create graph level embeddings that can be used to predict properties for the each graph. Remeber that the batch attribute will be essential for performining Global Pooling over our mini-batch of graphs.

[21]:

from ogb.graphproppred.mol_encoder import AtomEncoder
from torch_geometric.nn import global_add_pool, global_mean_pool

### GCN to predict graph property
class GCN_Graph(torch.nn.Module):
    def __init__(self, hidden_dim, output_dim, num_layers, dropout):
        super(GCN_Graph, self).__init__()

        # Load encoders for Atoms in molecule graphs
        self.node_encoder = AtomEncoder(hidden_dim)

        # Node embedding model
        # Note that the input_dim and output_dim are set to hidden_dim
        self.gnn_node = GCN(hidden_dim, hidden_dim,
            hidden_dim, num_layers, dropout, return_embeds=True)

        self.pool = None

        ############# Your code here ############
        ## Note:
        ## 1. Initialize self.pool as a global mean pooling layer
        ## For more information please refer to the documentation:
        ## https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#global-pooling-layers

        self.pool = global_mean_pool

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

        # Output layer
        self.linear = torch.nn.Linear(hidden_dim, output_dim)


    def reset_parameters(self):
      self.gnn_node.reset_parameters()
      self.linear.reset_parameters()

    def forward(self, batched_data):
        # TODO: Implement a function that takes as input a
        # mini-batch of graphs (torch_geometric.data.Batch) and
        # returns the predicted graph property for each graph.
        #
        # NOTE: Since we are predicting graph level properties,
        # your output will be a tensor with dimension equaling
        # the number of graphs in the mini-batch


        # Extract important attributes of our mini-batch
        x, edge_index, batch = batched_data.x, batched_data.edge_index, batched_data.batch
        embed = self.node_encoder(x)

        out = None

        ############# Your code here ############
        ## Note:
        ## 1. Construct node embeddings using existing GCN model
        ## 2. Use the global pooling layer to aggregate features for each individual graph
        ## For more information please refer to the documentation:
        ## https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#global-pooling-layers
        ## 3. Use a linear layer to predict each graph's property
        ## (~3 lines of code)

        embed = self.gnn_node(embed, edge_index)

        # 2. Use the global pooling layer
        #    to aggregate features for each individual graph
        out = self.pool(embed, batch)

        # 3. Use a linear layer to predict each graph's property
        out = self.linear(out)

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

        return out

[22]:

def train(model, device, data_loader, optimizer, loss_fn):
    # TODO: Implement a function that trains your model by
    # using the given optimizer and loss_fn.
    model.train()
    loss = 0

    for step, batch in enumerate(tqdm(data_loader, desc="Iteration")):
      batch = batch.to(device)

      if batch.x.shape[0] == 1 or batch.batch[-1] == 0:
          pass
      else:
        ## ignore nan targets (unlabeled) when computing training loss.
        is_labeled = batch.y == batch.y

        ############# Your code here ############
        ## Note:
        ## 1. Zero grad the optimizer
        ## 2. Feed the data into the model
        ## 3. Use `is_labeled` mask to filter output and labels
        ## 4. You may need to change the type of label to torch.float32
        ## 5. Feed the output and label to the loss_fn
        ## (~3 lines of code)

        # 1. Zero grad the optimizer
        optimizer.zero_grad()

        # 2. Feed the data into the model
        out = model(batch)

        # 3. Use `is_labeled` mask to filter output and labels
        loss = loss_fn(out[is_labeled], batch.y[is_labeled].float())

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

        loss.backward()
        optimizer.step()

    return loss.item()

[23]:

# The evaluation function
def eval(model, device, loader, evaluator, save_model_results=False, save_file=None):
    model.eval()
    y_true = []
    y_pred = []

    for step, batch in enumerate(tqdm(loader, desc="Iteration")):
        batch = batch.to(device)

        if batch.x.shape[0] == 1:
            pass
        else:
            with torch.no_grad():
                pred = model(batch)

            y_true.append(batch.y.view(pred.shape).detach().cpu())
            y_pred.append(pred.detach().cpu())

    y_true = torch.cat(y_true, dim = 0).numpy()
    y_pred = torch.cat(y_pred, dim = 0).numpy()

    input_dict = {"y_true": y_true, "y_pred": y_pred}

    if save_model_results:
        print ("Saving Model Predictions")

        # Create a pandas dataframe with a two columns
        # y_pred | y_true
        data = {}
        data['y_pred'] = y_pred.reshape(-1)
        data['y_true'] = y_true.reshape(-1)

        df = pd.DataFrame(data=data)
        # Save to csv
        df.to_csv('ogbg-molhiv_graph_' + save_file + '.csv', sep=',', index=False)

    return evaluator.eval(input_dict)

[24]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
  model = GCN_Graph(args['hidden_dim'],
              dataset.num_tasks, args['num_layers'],
              args['dropout']).to(device)
  evaluator = Evaluator(name='ogbg-molhiv')

[25]:

import copy

if 'IS_GRADESCOPE_ENV' not in os.environ:
  model.reset_parameters()

  optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])
  loss_fn = torch.nn.BCEWithLogitsLoss()

  best_model = None
  best_valid_acc = 0

  for epoch in range(1, 1 + args["epochs"]):
    print('Training...')
    loss = train(model, device, train_loader, optimizer, loss_fn)

    print('Evaluating...')
    train_result = eval(model, device, train_loader, evaluator)
    val_result = eval(model, device, valid_loader, evaluator)
    test_result = eval(model, device, test_loader, evaluator)

    train_acc, valid_acc, test_acc = train_result[dataset.eval_metric], val_result[dataset.eval_metric], test_result[dataset.eval_metric]
    if valid_acc > best_valid_acc:
        best_valid_acc = valid_acc
        best_model = copy.deepcopy(model)
    print(f'Epoch: {epoch:02d}, '
          f'Loss: {loss:.4f}, '
          f'Train: {100 * train_acc:.2f}%, '
          f'Valid: {100 * valid_acc:.2f}% '
          f'Test: {100 * test_acc:.2f}%')

Training...

Evaluating...

Epoch: 01, Loss: 0.0649, Train: 70.27%, Valid: 72.73% Test: 71.03%
Training...

Evaluating...

Epoch: 02, Loss: 0.0238, Train: 75.00%, Valid: 71.57% Test: 70.64%
Training...

Evaluating...

Epoch: 03, Loss: 0.0456, Train: 72.06%, Valid: 73.20% Test: 69.56%
Training...

Evaluating...

Epoch: 04, Loss: 0.0297, Train: 77.53%, Valid: 77.62% Test: 70.49%
Training...

Evaluating...

Epoch: 05, Loss: 0.0296, Train: 77.03%, Valid: 76.99% Test: 68.05%
Training...

Evaluating...

Epoch: 06, Loss: 0.0385, Train: 78.10%, Valid: 77.63% Test: 72.60%
Training...

Evaluating...

Epoch: 07, Loss: 0.0410, Train: 78.30%, Valid: 75.77% Test: 72.98%
Training...

Evaluating...

Epoch: 08, Loss: 0.0426, Train: 79.01%, Valid: 77.56% Test: 73.45%
Training...

Evaluating...

Epoch: 09, Loss: 0.0372, Train: 78.72%, Valid: 77.89% Test: 68.97%
Training...

Evaluating...

Epoch: 10, Loss: 0.0317, Train: 80.34%, Valid: 76.69% Test: 71.97%
Training...

Evaluating...

Epoch: 11, Loss: 0.0846, Train: 79.89%, Valid: 79.80% Test: 72.19%
Training...

Evaluating...

Epoch: 12, Loss: 0.0188, Train: 80.37%, Valid: 76.35% Test: 73.93%
Training...

Evaluating...

Epoch: 13, Loss: 0.0240, Train: 81.08%, Valid: 79.94% Test: 72.63%
Training...

Evaluating...

Epoch: 14, Loss: 0.0648, Train: 80.94%, Valid: 75.27% Test: 73.69%
Training...

Evaluating...

Epoch: 15, Loss: 0.0315, Train: 80.92%, Valid: 75.40% Test: 73.45%
Training...

Evaluating...

Epoch: 16, Loss: 0.5362, Train: 82.19%, Valid: 78.25% Test: 74.91%
Training...

Evaluating...

Epoch: 17, Loss: 0.0484, Train: 81.73%, Valid: 80.34% Test: 73.93%
Training...

Evaluating...

Epoch: 18, Loss: 0.0420, Train: 80.67%, Valid: 76.44% Test: 73.01%
Training...

Evaluating...

Epoch: 19, Loss: 0.0204, Train: 82.16%, Valid: 79.24% Test: 74.71%
Training...

Evaluating...

Epoch: 20, Loss: 0.0413, Train: 81.99%, Valid: 79.62% Test: 72.83%
Training...

Evaluating...

Epoch: 21, Loss: 0.0233, Train: 82.00%, Valid: 77.79% Test: 72.55%
Training...

Evaluating...

Epoch: 22, Loss: 0.0280, Train: 82.71%, Valid: 76.91% Test: 74.55%
Training...

Evaluating...

Epoch: 23, Loss: 0.0924, Train: 82.61%, Valid: 80.01% Test: 74.08%
Training...

Evaluating...

Epoch: 24, Loss: 0.0288, Train: 82.70%, Valid: 77.82% Test: 73.40%
Training...

Evaluating...

Epoch: 25, Loss: 0.5557, Train: 83.53%, Valid: 80.03% Test: 75.35%
Training...

Evaluating...

Epoch: 26, Loss: 0.2624, Train: 83.83%, Valid: 77.39% Test: 76.07%
Training...

Evaluating...

Epoch: 27, Loss: 0.0230, Train: 83.79%, Valid: 79.77% Test: 75.30%
Training...

Evaluating...

Epoch: 28, Loss: 0.0161, Train: 82.82%, Valid: 75.48% Test: 74.80%
Training...

Evaluating...

Epoch: 29, Loss: 0.0334, Train: 84.37%, Valid: 79.36% Test: 75.58%
Training...

Evaluating...

Epoch: 30, Loss: 0.0194, Train: 83.63%, Valid: 79.23% Test: 76.55%

Question 6: What are your `best_model` validation and test ROC-AUC scores? (20 points)¶

Run the cell below to see the results of your best of model and save your model’s predictions over the validation and test datasets. The resulting files are named ogbn-arxiv_graph_valid.csv and ogbn-arxiv_graph_test.csv.

Again, you can view these files by clicking on the Folder icon on the left side pannel. As in Colab 1, when you sumbit your assignment, you will have to download these files and attatch them to your submission.

[26]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
  train_acc = eval(best_model, device, train_loader, evaluator)[dataset.eval_metric]
  valid_acc = eval(best_model, device, valid_loader, evaluator, save_model_results=True, save_file="valid")[dataset.eval_metric]
  test_acc  = eval(best_model, device, test_loader, evaluator, save_model_results=True, save_file="test")[dataset.eval_metric]

  print(f'Best model: '
      f'Train: {100 * train_acc:.2f}%, '
      f'Valid: {100 * valid_acc:.2f}% '
      f'Test: {100 * test_acc:.2f}%')

Saving Model Predictions

Saving Model Predictions
Best model: Train: 81.73%, Valid: 80.34% Test: 73.93%

Question 7 (Optional): Experiment with the two other global pooling layers in Pytorch Geometric.¶

Submission¶

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

Your completed CS224W_Colab2.ipynb. From the “File” menu select “Download .ipynb” to save a local copy of your completed Colab. PLEASE DO NOT CHANGE THE NAME! The autograder depends on the .ipynb file being called “CS224W_Colab2.ipynb”.
ogbn-arxiv_node.csv
ogbg-molhiv_graph_valid.csv
ogbg-molhiv_graph_test.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-4 above and submit to gradescope. NOTE: DO NOT rename any of the downloaded files.

CS224W - Colab 2¶

Device¶

Setup¶

1) PyTorch Geometric (Datasets and Data)¶

PyG Datasets¶

Question 1: What is the number of classes and number of features in the ENZYMES dataset? (5 points)¶

PyG Data¶

Question 2: What is the label of the graph with index 100 in the ENZYMES dataset? (5 points)¶

Question 3: How many edges does the graph with index 200 have? (5 points)¶

2) Open Graph Benchmark (OGB)¶

Dataset and Data¶

Question 4: How many features are in the ogbn-arxiv graph? (5 points)¶

3) GNN: Node Property Prediction¶

Setup¶

Load and Preprocess the Dataset¶

GCN Model¶

Question 5: What are your best_model validation and test accuracies?(20 points)¶

4) GNN: Graph Property Prediction¶

Load and preprocess the dataset¶

Graph Prediction Model¶

Graph Mini-Batching¶

Implemention¶

Question 6: What are your best_model validation and test ROC-AUC scores? (20 points)¶

Question 7 (Optional): Experiment with the two other global pooling layers in Pytorch Geometric.¶

Submission¶

Question 5: What are your `best_model` validation and test accuracies?(20 points)¶

Question 6: What are your `best_model` validation and test ROC-AUC scores? (20 points)¶