Tutorial 9_1: Recurrent GNNs

In this tutorial we will implement an approximation of the Graph Neural Network Model (without enforcing contraction map) and analyze the GatedGraph Convolution of Pytorch Geometric.

import os
import torch
os.environ['TORCH'] = torch.__version__
print(torch.__version__)

!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git

1.12.1+cu113
     |████████████████████████████████| 7.9 MB 3.0 MB/s
     |████████████████████████████████| 3.5 MB 2.9 MB/s
  Building wheel for torch-geometric (setup.py) ... done

import os.path as osp
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_geometric.transforms as T
import torch_geometric
from torch_geometric.datasets import Planetoid, TUDataset
from torch_geometric.data import DataLoader
from torch_geometric.nn.inits import uniform
from torch.nn import Parameter as Param
from torch import Tensor
torch.manual_seed(42)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = "cpu"
from torch_geometric.nn.conv import MessagePassing

dataset = 'Cora'
transform = T.Compose([
    T.RandomNodeSplit('train_rest', num_val=500, num_test=500),
    T.TargetIndegree(),
])
path = osp.join('data', dataset)
dataset = Planetoid(path, dataset, transform=transform)
data = dataset[0]

Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.x
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.tx
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.allx
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.y
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ty
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ally
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.graph
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.test.index
Processing...
Done!

dataset = 'Cora'
path = osp.join('data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]
data = data.to(device)

Graph Neural Network Model

The MLP class is used to instantiate the transition and output functions as simple feed forard networks

class MLP(nn.Module):
    def __init__(self, input_dim, hid_dims, out_dim):
        super(MLP, self).__init__()

        self.mlp = nn.Sequential()
        dims = [input_dim] + hid_dims + [out_dim]
        for i in range(len(dims)-1):
            self.mlp.add_module('lay_{}'.format(i),nn.Linear(in_features=dims[i], out_features=dims[i+1]))
            if i+2 < len(dims):
                self.mlp.add_module('act_{}'.format(i), nn.Tanh())
    def reset_parameters(self):
        for i, l in enumerate(self.mlp):
            if type(l) == nn.Linear:
                nn.init.xavier_normal_(l.weight)

    def forward(self, x):
        return self.mlp(x)

The GNNM calss puts together the state propagations and the readout of the nodes’ states.

class GNNM(MessagePassing):
    def __init__(self, n_nodes, out_channels, features_dim, hid_dims, num_layers = 50, eps=1e-3, aggr = 'add',
                 bias = True, **kwargs):
        super(GNNM, self).__init__(aggr=aggr, **kwargs)

        self.node_states = Param(torch.zeros((n_nodes, features_dim)), requires_grad=False)
        self.out_channels = out_channels
        self.eps = eps
        self.num_layers = num_layers

        self.transition = MLP(features_dim, hid_dims, features_dim)
        self.readout = MLP(features_dim, hid_dims, out_channels)

        self.reset_parameters()
        print(self.transition)
        print(self.readout)

    def reset_parameters(self):
        self.transition.reset_parameters()
        self.readout.reset_parameters()

    def forward(self):
        edge_index = data.edge_index
        edge_weight = data.edge_attr
        node_states = self.node_states
        for i in range(self.num_layers):
            m = self.propagate(edge_index, x=node_states, edge_weight=edge_weight,
                               size=None)
            new_states = self.transition(m)
            with torch.no_grad():
                distance = torch.norm(new_states - node_states, dim=1)
                convergence = distance < self.eps
            node_states = new_states
            if convergence.all():
                break

        out = self.readout(node_states)

        return F.log_softmax(out, dim=-1)

    def message(self, x_j, edge_weight):
        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t, x) :
        return matmul(adj_t, x, reduce=self.aggr)

    def __repr__(self):
        return '{}({}, num_layers={})'.format(self.__class__.__name__,
                                              self.out_channels,
                                              self.num_layers)

model = GNNM(data.num_nodes, dataset.num_classes, 32, [64,64,64,64,64], eps=0.01).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()


test_dataset = dataset[:len(dataset) // 10]
train_dataset = dataset[len(dataset) // 10:]
test_loader = DataLoader(test_dataset)
train_loader = DataLoader(train_dataset)

def train():
    model.train()
    optimizer.zero_grad()
    loss_fn(model()[data.train_mask], data.y[data.train_mask]).backward()
    optimizer.step()


def test():
    model.eval()
    logits, accs = model(), []
    for _, mask in data('train_mask', 'val_mask', 'test_mask'):
        pred = logits[mask].max(1)[1]
        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs


for epoch in range(1, 51):
    train()
    accs = test()
    train_acc = accs[0]
    val_acc = accs[1]
    test_acc = accs[2]
    print('Epoch: {:03d}, Train Acc: {:.5f}, '
          'Val Acc: {:.5f}, Test Acc: {:.5f}'.format(epoch, train_acc,
                                                       val_acc, test_acc))

/usr/local/lib/python3.7/dist-packages/torch_geometric/deprecation.py:12: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead
  warnings.warn(out)

MLP(
  (mlp): Sequential(
    (lay_0): Linear(in_features=32, out_features=64, bias=True)
    (act_0): Tanh()
    (lay_1): Linear(in_features=64, out_features=64, bias=True)
    (act_1): Tanh()
    (lay_2): Linear(in_features=64, out_features=64, bias=True)
    (act_2): Tanh()
    (lay_3): Linear(in_features=64, out_features=64, bias=True)
    (act_3): Tanh()
    (lay_4): Linear(in_features=64, out_features=64, bias=True)
    (act_4): Tanh()
    (lay_5): Linear(in_features=64, out_features=32, bias=True)
  )
)
MLP(
  (mlp): Sequential(
    (lay_0): Linear(in_features=32, out_features=64, bias=True)
    (act_0): Tanh()
    (lay_1): Linear(in_features=64, out_features=64, bias=True)
    (act_1): Tanh()
    (lay_2): Linear(in_features=64, out_features=64, bias=True)
    (act_2): Tanh()
    (lay_3): Linear(in_features=64, out_features=64, bias=True)
    (act_3): Tanh()
    (lay_4): Linear(in_features=64, out_features=64, bias=True)
    (act_4): Tanh()
    (lay_5): Linear(in_features=64, out_features=7, bias=True)
  )
)
Epoch: 001, Train Acc: 0.12857, Val Acc: 0.06800, Test Acc: 0.08800
Epoch: 002, Train Acc: 0.14286, Val Acc: 0.25200, Test Acc: 0.25100
Epoch: 003, Train Acc: 0.12143, Val Acc: 0.24400, Test Acc: 0.26100
Epoch: 004, Train Acc: 0.17143, Val Acc: 0.20200, Test Acc: 0.20000
Epoch: 005, Train Acc: 0.16429, Val Acc: 0.23000, Test Acc: 0.22900
Epoch: 006, Train Acc: 0.22143, Val Acc: 0.09800, Test Acc: 0.10400
Epoch: 007, Train Acc: 0.14286, Val Acc: 0.11400, Test Acc: 0.09900
Epoch: 008, Train Acc: 0.14286, Val Acc: 0.10000, Test Acc: 0.08300
Epoch: 009, Train Acc: 0.12857, Val Acc: 0.08800, Test Acc: 0.08100
Epoch: 010, Train Acc: 0.12857, Val Acc: 0.15400, Test Acc: 0.14000
Epoch: 011, Train Acc: 0.14286, Val Acc: 0.12400, Test Acc: 0.11600
Epoch: 012, Train Acc: 0.14286, Val Acc: 0.07400, Test Acc: 0.09400
Epoch: 013, Train Acc: 0.13571, Val Acc: 0.07800, Test Acc: 0.08600
Epoch: 014, Train Acc: 0.15000, Val Acc: 0.11600, Test Acc: 0.11600
Epoch: 015, Train Acc: 0.19286, Val Acc: 0.20400, Test Acc: 0.19000
Epoch: 016, Train Acc: 0.15714, Val Acc: 0.16600, Test Acc: 0.15400
Epoch: 017, Train Acc: 0.17143, Val Acc: 0.15600, Test Acc: 0.13700
Epoch: 018, Train Acc: 0.18571, Val Acc: 0.14000, Test Acc: 0.11900
Epoch: 019, Train Acc: 0.21429, Val Acc: 0.14000, Test Acc: 0.11200
Epoch: 020, Train Acc: 0.18571, Val Acc: 0.13400, Test Acc: 0.12000
Epoch: 021, Train Acc: 0.18571, Val Acc: 0.15000, Test Acc: 0.12900
Epoch: 022, Train Acc: 0.20000, Val Acc: 0.17000, Test Acc: 0.14000
Epoch: 023, Train Acc: 0.16429, Val Acc: 0.11800, Test Acc: 0.09800
Epoch: 024, Train Acc: 0.17143, Val Acc: 0.13400, Test Acc: 0.11300
Epoch: 025, Train Acc: 0.20714, Val Acc: 0.14200, Test Acc: 0.12200
Epoch: 026, Train Acc: 0.16429, Val Acc: 0.16200, Test Acc: 0.14700
Epoch: 027, Train Acc: 0.20714, Val Acc: 0.14000, Test Acc: 0.11100
Epoch: 028, Train Acc: 0.21429, Val Acc: 0.12800, Test Acc: 0.10300
Epoch: 029, Train Acc: 0.20000, Val Acc: 0.10600, Test Acc: 0.10200
Epoch: 030, Train Acc: 0.22143, Val Acc: 0.12000, Test Acc: 0.11000
Epoch: 031, Train Acc: 0.22857, Val Acc: 0.13200, Test Acc: 0.11900
Epoch: 032, Train Acc: 0.20714, Val Acc: 0.13600, Test Acc: 0.12400
Epoch: 033, Train Acc: 0.22143, Val Acc: 0.14800, Test Acc: 0.12700
Epoch: 034, Train Acc: 0.21429, Val Acc: 0.18200, Test Acc: 0.13200
Epoch: 035, Train Acc: 0.20714, Val Acc: 0.17400, Test Acc: 0.13400
Epoch: 036, Train Acc: 0.20714, Val Acc: 0.18400, Test Acc: 0.14400
Epoch: 037, Train Acc: 0.20714, Val Acc: 0.19400, Test Acc: 0.14400
Epoch: 038, Train Acc: 0.21429, Val Acc: 0.19800, Test Acc: 0.14800
Epoch: 039, Train Acc: 0.24286, Val Acc: 0.16200, Test Acc: 0.14900
Epoch: 040, Train Acc: 0.25714, Val Acc: 0.15800, Test Acc: 0.14100
Epoch: 041, Train Acc: 0.24286, Val Acc: 0.15400, Test Acc: 0.14100
Epoch: 042, Train Acc: 0.24286, Val Acc: 0.14800, Test Acc: 0.14200
Epoch: 043, Train Acc: 0.23571, Val Acc: 0.14600, Test Acc: 0.13600
Epoch: 044, Train Acc: 0.22143, Val Acc: 0.15200, Test Acc: 0.13500
Epoch: 045, Train Acc: 0.23571, Val Acc: 0.14600, Test Acc: 0.14000
Epoch: 046, Train Acc: 0.28571, Val Acc: 0.12600, Test Acc: 0.13000
Epoch: 047, Train Acc: 0.20714, Val Acc: 0.17000, Test Acc: 0.13900
Epoch: 048, Train Acc: 0.27857, Val Acc: 0.18800, Test Acc: 0.14200
Epoch: 049, Train Acc: 0.14286, Val Acc: 0.05800, Test Acc: 0.06400
Epoch: 050, Train Acc: 0.17143, Val Acc: 0.11000, Test Acc: 0.09300

Gated Graph Neural Network

class GatedGraphConv(MessagePassing):

    def __init__(self, out_channels, num_layers, aggr = 'add',
                 bias = True, **kwargs):
        super(GatedGraphConv, self).__init__(aggr=aggr, **kwargs)

        self.out_channels = out_channels
        self.num_layers = num_layers

        self.weight = Param(Tensor(num_layers, out_channels, out_channels))
        self.rnn = torch.nn.GRUCell(out_channels, out_channels, bias=bias)

        self.reset_parameters()

    def reset_parameters(self):
        uniform(self.out_channels, self.weight)
        self.rnn.reset_parameters()

    def forward(self, data):
        """"""
        x = data.x
        edge_index = data.edge_index
        edge_weight = data.edge_attr
        if x.size(-1) > self.out_channels:
            raise ValueError('The number of input channels is not allowed to '
                             'be larger than the number of output channels')

        if x.size(-1) < self.out_channels:
            zero = x.new_zeros(x.size(0), self.out_channels - x.size(-1))
            x = torch.cat([x, zero], dim=1)

        for i in range(self.num_layers):
            m = torch.matmul(x, self.weight[i])
            m = self.propagate(edge_index, x=m, edge_weight=edge_weight,
                               size=None)
            x = self.rnn(m, x)

        return x

    def message(self, x_j, edge_weight):
        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t, x):
        return matmul(adj_t, x, reduce=self.aggr)

    def __repr__(self):
        return '{}({}, num_layers={})'.format(self.__class__.__name__,
                                              self.out_channels,
                                              self.num_layers)

class GGNN(torch.nn.Module):
    def __init__(self):
        super(GGNN, self).__init__()

        self.conv = GatedGraphConv(1433, 3)
        self.mlp = MLP(1433, [32,32,32], dataset.num_classes)

    def forward(self):
        x = self.conv(data)
        x = self.mlp(x)
        return F.log_softmax(x, dim=-1)

device = "cpu"
model = GGNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()


test_dataset = dataset[:len(dataset) // 10]
train_dataset = dataset[len(dataset) // 10:]
test_loader = DataLoader(test_dataset)
train_loader = DataLoader(train_dataset)

def train():
    model.train()
    optimizer.zero_grad()
    loss_fn(model()[data.train_mask], data.y[data.train_mask]).backward()
    optimizer.step()


def test():
    model.eval()
    logits, accs = model(), []
    for _, mask in data('train_mask', 'val_mask', 'test_mask'):
        pred = logits[mask].max(1)[1]
        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs


for epoch in range(1, 51):
    train()
    accs = test()
    train_acc = accs[0]
    val_acc = accs[1]
    test_acc = accs[2]
    print('Epoch: {:03d}, Train Acc: {:.5f}, '
          'Val Acc: {:.5f}, Test Acc: {:.5f}'.format(epoch, train_acc,
                                                       val_acc, test_acc))

Epoch: 001, Train Acc: 0.27143, Val Acc: 0.15800, Test Acc: 0.15400
Epoch: 002, Train Acc: 0.35000, Val Acc: 0.22200, Test Acc: 0.22200
Epoch: 003, Train Acc: 0.18571, Val Acc: 0.22400, Test Acc: 0.21000
Epoch: 004, Train Acc: 0.32857, Val Acc: 0.29600, Test Acc: 0.28300
Epoch: 005, Train Acc: 0.36429, Val Acc: 0.25000, Test Acc: 0.26700
Epoch: 006, Train Acc: 0.45000, Val Acc: 0.35000, Test Acc: 0.37200
Epoch: 007, Train Acc: 0.54286, Val Acc: 0.45800, Test Acc: 0.47100
Epoch: 008, Train Acc: 0.53571, Val Acc: 0.40800, Test Acc: 0.43000
Epoch: 009, Train Acc: 0.60000, Val Acc: 0.51000, Test Acc: 0.50700
Epoch: 010, Train Acc: 0.64286, Val Acc: 0.58800, Test Acc: 0.57500
Epoch: 011, Train Acc: 0.62143, Val Acc: 0.57600, Test Acc: 0.57700
Epoch: 012, Train Acc: 0.63571, Val Acc: 0.54000, Test Acc: 0.55100
Epoch: 013, Train Acc: 0.65714, Val Acc: 0.55000, Test Acc: 0.55200
Epoch: 014, Train Acc: 0.67143, Val Acc: 0.59200, Test Acc: 0.57100
Epoch: 015, Train Acc: 0.70714, Val Acc: 0.60800, Test Acc: 0.58400
Epoch: 016, Train Acc: 0.73571, Val Acc: 0.62400, Test Acc: 0.60400
Epoch: 017, Train Acc: 0.74286, Val Acc: 0.61200, Test Acc: 0.58800
Epoch: 018, Train Acc: 0.69286, Val Acc: 0.57000, Test Acc: 0.55700
Epoch: 019, Train Acc: 0.70714, Val Acc: 0.59400, Test Acc: 0.58700
Epoch: 020, Train Acc: 0.74286, Val Acc: 0.59600, Test Acc: 0.60400
Epoch: 021, Train Acc: 0.72143, Val Acc: 0.58600, Test Acc: 0.59000
Epoch: 022, Train Acc: 0.71429, Val Acc: 0.57800, Test Acc: 0.56100
Epoch: 023, Train Acc: 0.72857, Val Acc: 0.56800, Test Acc: 0.55100
Epoch: 024, Train Acc: 0.74286, Val Acc: 0.58200, Test Acc: 0.57100
Epoch: 025, Train Acc: 0.82143, Val Acc: 0.59600, Test Acc: 0.61400
Epoch: 026, Train Acc: 0.83571, Val Acc: 0.58800, Test Acc: 0.59700
Epoch: 027, Train Acc: 0.83571, Val Acc: 0.57800, Test Acc: 0.58000
Epoch: 028, Train Acc: 0.83571, Val Acc: 0.57800, Test Acc: 0.57800
Epoch: 029, Train Acc: 0.82143, Val Acc: 0.57800, Test Acc: 0.58200
Epoch: 030, Train Acc: 0.82857, Val Acc: 0.58600, Test Acc: 0.57900
Epoch: 031, Train Acc: 0.83571, Val Acc: 0.59200, Test Acc: 0.58600
Epoch: 032, Train Acc: 0.86429, Val Acc: 0.61000, Test Acc: 0.59700
Epoch: 033, Train Acc: 0.94286, Val Acc: 0.62000, Test Acc: 0.62300
Epoch: 034, Train Acc: 0.95714, Val Acc: 0.62000, Test Acc: 0.62900
Epoch: 035, Train Acc: 0.97143, Val Acc: 0.62000, Test Acc: 0.63300
Epoch: 036, Train Acc: 0.98571, Val Acc: 0.61400, Test Acc: 0.63400
Epoch: 037, Train Acc: 0.98571, Val Acc: 0.60800, Test Acc: 0.63600
Epoch: 038, Train Acc: 0.98571, Val Acc: 0.60800, Test Acc: 0.63900
Epoch: 039, Train Acc: 0.98571, Val Acc: 0.61400, Test Acc: 0.63100
Epoch: 040, Train Acc: 0.98571, Val Acc: 0.61600, Test Acc: 0.62900
Epoch: 041, Train Acc: 0.98571, Val Acc: 0.61800, Test Acc: 0.63200
Epoch: 042, Train Acc: 0.98571, Val Acc: 0.61000, Test Acc: 0.63200
Epoch: 043, Train Acc: 0.98571, Val Acc: 0.60600, Test Acc: 0.63200
Epoch: 044, Train Acc: 0.98571, Val Acc: 0.60400, Test Acc: 0.63100
Epoch: 045, Train Acc: 0.98571, Val Acc: 0.60600, Test Acc: 0.63100
Epoch: 046, Train Acc: 0.98571, Val Acc: 0.60400, Test Acc: 0.62900
Epoch: 047, Train Acc: 0.98571, Val Acc: 0.60400, Test Acc: 0.63100
Epoch: 048, Train Acc: 0.98571, Val Acc: 0.60400, Test Acc: 0.63300
Epoch: 049, Train Acc: 0.98571, Val Acc: 0.60400, Test Acc: 0.63400
Epoch: 050, Train Acc: 0.99286, Val Acc: 0.60600, Test Acc: 0.62800