Solution to Colab problem (Graph Classification) using NNConv

[adhikarirsr]

We can get a decent result for graph classification problem using GCNConv, GraphConv and GATConv. It also seems like setting RandomSeed to 12345 helps to get around 82% accuracy.

Since we have access to edge_attr, I wanted to use a GNN that can leverage edge_attr. Here's what I did:

#NNConv
from torch_geometric.nn import NNConv
import torch.nn as nn

class Net(torch.nn.Module):
    def __init__(self,hidden_channels):
        super(Net, self).__init__()
        nn1 = nn.Sequential(nn.Linear(4, 25), nn.ReLU(), nn.Linear(25, 32))
        self.conv1 = NNConv(dataset.num_node_features, hidden_channels, nn1, aggr='add')

        nn2 = nn.Sequential(nn.Linear(4, 25), nn.ReLU(), nn.Linear(25, 1024))
        self.conv2 = NNConv(hidden_channels, hidden_channels, nn2, aggr='add')
        
        self.lin = Linear(hidden_channels, dataset.num_classes)

        
    def forward(self, x, edge_index, edge_attr, batch):
        x = F.relu(self.conv1(x, edge_index, edge_attr))
        x = F.relu(self.conv2(x, edge_index, edge_attr))
                
        x = global_mean_pool(x, batch)
        
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.lin(x)
        return x

model = Net(hidden_channels=32)
print(model)

from IPython.display import Javascript
display(Javascript('''google.colab.output.setIframeHeight(0, true, {maxHeight: 300})'''))

model = Net(hidden_channels=32)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()

def train():
    model.train()

    for data in train_loader:  # Iterate in batches over the training dataset.
         out = model(data.x, data.edge_index, data.edge_attr, data.batch)  # Perform a single forward pass.
         loss = criterion(out, data.y)  # Compute the loss.
         loss.backward()  # Derive gradients.
         optimizer.step()  # Update parameters based on gradients.
         optimizer.zero_grad()  # Clear gradients.

def test(loader):
     model.eval()

     correct = 0
     for data in loader:  # Iterate in batches over the training/test dataset.
         out = model(data.x, data.edge_index, data.edge_attr,data.batch)  
         pred = out.argmax(dim=1)  # Use the class with highest probability.
         correct += int((pred == data.y).sum())  # Check against ground-truth labels.
     return correct / len(loader.dataset)  # Derive ratio of correct predictions.


for epoch in range(1, 201):
    train()
    train_acc = test(train_loader)
    test_acc = test(test_loader)
    print(f'Epoch: {epoch:03d}, Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')

I am getting this error:

/usr/local/lib/python3.7/dist-packages/torch_geometric/nn/conv/nn_conv.py in message(self, x_j, edge_attr)
    100     def message(self, x_j: Tensor, edge_attr: Tensor) -> Tensor:
    101         weight = self.nn(edge_attr)
--> 102         weight = weight.view(-1, self.in_channels_l, self.out_channels)
    103         return torch.matmul(x_j.unsqueeze(1), weight).squeeze(1)
    104 

RuntimeError: shape '[-1, 7, 32]' is invalid for input of size 77696

Why am I getting this error?

[rusty1s]

Your first nn needs to be initialized as:

nn1 = nn.Sequential(nn.Linear(4, 25), nn.ReLU(), nn.Linear(25, 7 * 32))

[adhikarirsr]

Thanks @rusty1s. I made such a silly mistake.
Interestingly enough, it does not do better than GraphConv and GATConv. So, it seems like GATConv learns a better(or as good) edge attributes than the one that already exist in the graph?

[rusty1s]

Yes, that can happen, in particular if your edge features are not that important for the given task. It also introduces more parameters, and you might see over-fitting.