Tutorial 9_2: 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 5.5 MB/s
     |████████████████████████████████| 3.5 MB 5.4 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.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.18571, Val Acc: 0.21000, Test Acc: 0.18800
Epoch: 002, Train Acc: 0.15000, Val Acc: 0.16600, Test Acc: 0.14400
Epoch: 003, Train Acc: 0.19286, Val Acc: 0.13400, Test Acc: 0.11300
Epoch: 004, Train Acc: 0.15714, Val Acc: 0.12800, Test Acc: 0.09400
Epoch: 005, Train Acc: 0.10714, Val Acc: 0.12600, Test Acc: 0.12600
Epoch: 006, Train Acc: 0.17857, Val Acc: 0.21200, Test Acc: 0.20400
Epoch: 007, Train Acc: 0.17857, Val Acc: 0.20600, Test Acc: 0.20000
Epoch: 008, Train Acc: 0.17857, Val Acc: 0.19000, Test Acc: 0.19700
Epoch: 009, Train Acc: 0.22857, Val Acc: 0.17400, Test Acc: 0.17100
Epoch: 010, Train Acc: 0.21429, Val Acc: 0.12800, Test Acc: 0.11000
Epoch: 011, Train Acc: 0.20714, Val Acc: 0.14000, Test Acc: 0.11500
Epoch: 012, Train Acc: 0.20714, Val Acc: 0.15400, Test Acc: 0.12700
Epoch: 013, Train Acc: 0.20714, Val Acc: 0.16800, Test Acc: 0.13000
Epoch: 014, Train Acc: 0.21429, Val Acc: 0.15800, Test Acc: 0.12600
Epoch: 015, Train Acc: 0.20000, Val Acc: 0.17200, Test Acc: 0.14000
Epoch: 016, Train Acc: 0.24286, Val Acc: 0.18000, Test Acc: 0.15600
Epoch: 017, Train Acc: 0.26429, Val Acc: 0.17400, Test Acc: 0.16000
Epoch: 018, Train Acc: 0.25714, Val Acc: 0.19200, Test Acc: 0.15900
Epoch: 019, Train Acc: 0.25714, Val Acc: 0.18000, Test Acc: 0.15800
Epoch: 020, Train Acc: 0.24286, Val Acc: 0.17800, Test Acc: 0.14700
Epoch: 021, Train Acc: 0.24286, Val Acc: 0.15200, Test Acc: 0.13600
Epoch: 022, Train Acc: 0.23571, Val Acc: 0.16600, Test Acc: 0.12500
Epoch: 023, Train Acc: 0.25000, Val Acc: 0.17400, Test Acc: 0.12900
Epoch: 024, Train Acc: 0.25000, Val Acc: 0.17800, Test Acc: 0.13700
Epoch: 025, Train Acc: 0.25000, Val Acc: 0.18200, Test Acc: 0.14000
Epoch: 026, Train Acc: 0.25000, Val Acc: 0.19200, Test Acc: 0.15200
Epoch: 027, Train Acc: 0.27143, Val Acc: 0.18400, Test Acc: 0.15800
Epoch: 028, Train Acc: 0.25000, Val Acc: 0.18600, Test Acc: 0.14700
Epoch: 029, Train Acc: 0.25714, Val Acc: 0.18400, Test Acc: 0.14700
Epoch: 030, Train Acc: 0.25000, Val Acc: 0.17600, Test Acc: 0.14700
Epoch: 031, Train Acc: 0.23571, Val Acc: 0.17000, Test Acc: 0.14100
Epoch: 032, Train Acc: 0.23571, Val Acc: 0.17000, Test Acc: 0.14100
Epoch: 033, Train Acc: 0.25714, Val Acc: 0.18200, Test Acc: 0.14800
Epoch: 034, Train Acc: 0.26429, Val Acc: 0.19000, Test Acc: 0.15400
Epoch: 035, Train Acc: 0.27143, Val Acc: 0.19000, Test Acc: 0.15500
Epoch: 036, Train Acc: 0.25000, Val Acc: 0.19600, Test Acc: 0.15000
Epoch: 037, Train Acc: 0.27143, Val Acc: 0.17800, Test Acc: 0.15400
Epoch: 038, Train Acc: 0.22857, Val Acc: 0.17400, Test Acc: 0.13200
Epoch: 039, Train Acc: 0.24286, Val Acc: 0.16800, Test Acc: 0.14900
Epoch: 040, Train Acc: 0.25714, Val Acc: 0.18800, Test Acc: 0.15200
Epoch: 041, Train Acc: 0.22143, Val Acc: 0.17200, Test Acc: 0.13900
Epoch: 042, Train Acc: 0.25714, Val Acc: 0.17400, Test Acc: 0.15300
Epoch: 043, Train Acc: 0.22857, Val Acc: 0.18200, Test Acc: 0.15500
Epoch: 044, Train Acc: 0.23571, Val Acc: 0.17600, Test Acc: 0.14300
Epoch: 045, Train Acc: 0.27857, Val Acc: 0.18600, Test Acc: 0.15900
Epoch: 046, Train Acc: 0.23571, Val Acc: 0.17600, Test Acc: 0.15300
Epoch: 047, Train Acc: 0.25000, Val Acc: 0.18200, Test Acc: 0.15200
Epoch: 048, Train Acc: 0.26429, Val Acc: 0.18200, Test Acc: 0.15900
Epoch: 049, Train Acc: 0.23571, Val Acc: 0.18600, Test Acc: 0.14400
Epoch: 050, Train Acc: 0.28571, Val Acc: 0.17000, Test Acc: 0.15000

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)

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.15000, Val Acc: 0.13600, Test Acc: 0.13100
Epoch: 002, Train Acc: 0.17857, Val Acc: 0.25400, Test Acc: 0.23500
Epoch: 003, Train Acc: 0.18571, Val Acc: 0.13400, Test Acc: 0.13800
Epoch: 004, Train Acc: 0.29286, Val Acc: 0.22000, Test Acc: 0.22500
Epoch: 005, Train Acc: 0.36429, Val Acc: 0.34800, Test Acc: 0.33200
Epoch: 006, Train Acc: 0.40000, Val Acc: 0.36800, Test Acc: 0.34100
Epoch: 007, Train Acc: 0.40714, Val Acc: 0.37800, Test Acc: 0.35300
Epoch: 008, Train Acc: 0.42857, Val Acc: 0.41800, Test Acc: 0.39000
Epoch: 009, Train Acc: 0.54286, Val Acc: 0.48800, Test Acc: 0.46900
Epoch: 010, Train Acc: 0.55714, Val Acc: 0.51800, Test Acc: 0.50300
Epoch: 011, Train Acc: 0.50714, Val Acc: 0.48400, Test Acc: 0.47100
Epoch: 012, Train Acc: 0.46429, Val Acc: 0.46400, Test Acc: 0.46000
Epoch: 013, Train Acc: 0.50000, Val Acc: 0.48600, Test Acc: 0.49000
Epoch: 014, Train Acc: 0.57857, Val Acc: 0.51600, Test Acc: 0.51200
Epoch: 015, Train Acc: 0.60714, Val Acc: 0.54000, Test Acc: 0.53400
Epoch: 016, Train Acc: 0.68571, Val Acc: 0.53600, Test Acc: 0.54000
Epoch: 017, Train Acc: 0.69286, Val Acc: 0.54600, Test Acc: 0.55200
Epoch: 018, Train Acc: 0.66429, Val Acc: 0.55000, Test Acc: 0.55000
Epoch: 019, Train Acc: 0.70000, Val Acc: 0.54000, Test Acc: 0.55100
Epoch: 020, Train Acc: 0.75714, Val Acc: 0.53000, Test Acc: 0.54800
Epoch: 021, Train Acc: 0.77857, Val Acc: 0.57000, Test Acc: 0.58700
Epoch: 022, Train Acc: 0.78571, Val Acc: 0.57200, Test Acc: 0.60900
Epoch: 023, Train Acc: 0.80000, Val Acc: 0.57400, Test Acc: 0.61600
Epoch: 024, Train Acc: 0.80000, Val Acc: 0.56400, Test Acc: 0.61300
Epoch: 025, Train Acc: 0.81429, Val Acc: 0.57000, Test Acc: 0.61200
Epoch: 026, Train Acc: 0.80714, Val Acc: 0.54800, Test Acc: 0.57600
Epoch: 027, Train Acc: 0.81429, Val Acc: 0.56200, Test Acc: 0.57400
Epoch: 028, Train Acc: 0.82143, Val Acc: 0.56200, Test Acc: 0.58200
Epoch: 029, Train Acc: 0.82857, Val Acc: 0.55800, Test Acc: 0.58700
Epoch: 030, Train Acc: 0.82857, Val Acc: 0.55600, Test Acc: 0.58900
Epoch: 031, Train Acc: 0.82857, Val Acc: 0.55400, Test Acc: 0.58400
Epoch: 032, Train Acc: 0.91429, Val Acc: 0.58400, Test Acc: 0.61700
Epoch: 033, Train Acc: 0.92857, Val Acc: 0.59400, Test Acc: 0.64300
Epoch: 034, Train Acc: 0.94286, Val Acc: 0.59800, Test Acc: 0.64500
Epoch: 035, Train Acc: 0.94286, Val Acc: 0.60200, Test Acc: 0.64600
Epoch: 036, Train Acc: 0.95714, Val Acc: 0.62400, Test Acc: 0.66000
Epoch: 037, Train Acc: 0.95714, Val Acc: 0.62800, Test Acc: 0.66300
Epoch: 038, Train Acc: 0.96429, Val Acc: 0.64000, Test Acc: 0.67100
Epoch: 039, Train Acc: 0.96429, Val Acc: 0.64600, Test Acc: 0.67600
Epoch: 040, Train Acc: 0.96429, Val Acc: 0.64000, Test Acc: 0.67200
Epoch: 041, Train Acc: 0.95714, Val Acc: 0.64800, Test Acc: 0.68500
Epoch: 042, Train Acc: 0.96429, Val Acc: 0.64200, Test Acc: 0.67100
Epoch: 043, Train Acc: 0.96429, Val Acc: 0.64800, Test Acc: 0.66200
Epoch: 044, Train Acc: 0.84286, Val Acc: 0.64400, Test Acc: 0.66100
Epoch: 045, Train Acc: 0.95000, Val Acc: 0.64400, Test Acc: 0.65400
Epoch: 046, Train Acc: 0.82143, Val Acc: 0.64200, Test Acc: 0.64100
Epoch: 047, Train Acc: 0.92143, Val Acc: 0.64800, Test Acc: 0.64800
Epoch: 048, Train Acc: 0.94286, Val Acc: 0.63200, Test Acc: 0.63200
Epoch: 049, Train Acc: 0.95714, Val Acc: 0.64000, Test Acc: 0.63800
Epoch: 050, Train Acc: 0.92143, Val Acc: 0.64400, Test Acc: 0.64500