GAT for graph classification

[dokato]

I'm having trouble of reimplementing a GAT Network from this tutorial on the PyTorch geometric website: https://colab.research.google.com/github/AntonioLonga/PytorchGeometricTutorial/blob/main/Tutorial3/Tutorial3.ipynb

...to work on a graph classification problem. For instance, let's take a MNISTSuperpixels dataset. How to make it work so it trains and predict my data?

Here's is my first try that fails miserably:

import os

import numpy as np
import torch

import torch.nn as nn
import torch.nn.functional as F

from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GATConv
from torch_geometric.datasets import Planetoid
from torch_geometric.datasets import MNISTSuperpixels

import torch_geometric.transforms as T

import matplotlib.pyplot as plt

batch_size = 32

path = os.path.join(os.path.dirname(os.getcwd()), 'data', 'MNIST')
name_data = 'MNISTSuperpix'
dataset =  MNISTSuperpixels(path, True, transform=T.Cartesian())
dataset.transform = T.NormalizeFeatures()

train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

print(f"Number of Classes in {name_data}:", dataset.num_classes)
print(f"Number of Node Features in {name_data}:", dataset.num_node_features)


class GAT(torch.nn.Module):
    def __init__(self):
        super(GAT, self).__init__()
        self.hid = 8
        self.in_head = 8
        self.out_head = 1
        
        self.conv1 = GATConv(dataset.num_features, self.hid, heads=self.in_head, dropout=0.6)
        self.conv2 = GATConv(self.hid*self.in_head, dataset.num_classes, concat=False,
                             heads=self.out_head, dropout=0.6)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
                
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv1(x, edge_index)
        x = F.elu(x)
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        
        return F.log_softmax(x, dim=1)
    
    
    
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device', device)

model = GAT().to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

model.train()
for epoch in range(100):
    model.train()
    loss_all = 0
    for batch_data in train_loader:
        batch_data = batch_data.to(device)
        optimizer.zero_grad()
        out = model(batch_data)
        loss = F.nll_loss(out, batch_data.y)
        loss.backward()
        loss_all += batch_data.num_graphs * loss.item()
        optimizer.step()
    if epoch%10 == 0:
        print(f'epoch: {epoch}, loss: {loss_all}')

Currently, I get some error related to the input batch size...

[rusty1s]

Note that MNIST is a graph-level task, so you need to convert your node-level features into graph-level ones, e.g. via global_mean_pool. Here is my updated example:

import os

import numpy as np
import torch

import torch.nn as nn
import torch.nn.functional as F

from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GATConv, global_mean_pool
from torch_geometric.datasets import Planetoid
from torch_geometric.datasets import MNISTSuperpixels

import torch_geometric.transforms as T

import matplotlib.pyplot as plt

batch_size = 32

path = os.path.join(os.path.dirname(os.getcwd()), 'data', 'MNIST')
name_data = 'MNISTSuperpix'
dataset = MNISTSuperpixels(path, True, transform=T.Cartesian())
dataset.transform = T.NormalizeFeatures()

train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

print(f"Number of Classes in {name_data}:", dataset.num_classes)
print(f"Number of Node Features in {name_data}:", dataset.num_node_features)


class GAT(torch.nn.Module):
    def __init__(self):
        super(GAT, self).__init__()
        self.hid = 8
        self.head = 8

        self.conv1 = GATConv(3, self.hid, heads=self.head, dropout=0.6)
        self.conv2 = GATConv(self.hid * self.head, self.hid, heads=self.head,
                             dropout=0.6)

        self.lin = nn.Linear(self.hid * self.head, dataset.num_classes)

    def forward(self, data):
        x, pos, edge_index = data.x, data.pos, data.edge_index,

        x = torch.cat([x, pos], dim=-1)

        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv1(x, edge_index)
        x = F.elu(x)
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        x = F.elu(x)

        x = global_mean_pool(x, data.batch)
        x = self.lin(x)

        return F.log_softmax(x, dim=1)


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

model = GAT().to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

model.train()
for epoch in range(100):
    model.train()
    loss_all = 0
    for batch_data in train_loader:
        batch_data = batch_data.to(device)
        optimizer.zero_grad()
        out = model(batch_data)
        loss = F.nll_loss(out, batch_data.y)
        loss.backward()
        loss_all += batch_data.num_graphs * loss.item()
        optimizer.step()
    print(f'epoch: {epoch}, loss: {loss_all}')

Furthermore, I added the node positions data.pos to the initial node features.

[dokato]

Ah, of course, thanks a lot @rusty1s !