AssertionError with example for heterogeneous GNN

[Zytrus]

Hello,

I have encountered a mysterios error while trying to execute a graph classification with one of the examples given in Heterogeneous GNN - Automatically Converting GNN Models (GAT).
I specifically tried out the GAT GNN:

from torch_geometric.nn import GATConv, Linear, to_hetero

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x


model = GAT(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

and am training it like this:

def train():
    model.train()

    for data_train in train_dataset:  # Iterate in batches over the training dataset.
        out = model(data_train.x_dict, data_train.edge_index_dict)  # Perform a single forward pass.
        loss = criterion(out, data_train.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_test in loader:  # Iterate in batches over the training/test dataset.
        out = model(data_test.x_dict, data_test.edge_index_dict)  
        pred = out.argmax(dim=1)  # Use the class with highest probability.
        correct += int((pred == data_test.y).sum())  # Check against ground-truth labels.
    return correct / len(loader)  # Derive ratio of correct predictions.

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

Since my dataset consists of only 29 graphs, I don't need to batch the graphs, but I get this mysterious error:

Traceback (most recent call last):
  File "C:\Users\Documents\Privacy\test.py", line 89, in <module>
    train()
  File "C:\Users\Documents\Privacy\test.py", line 71, in train
    out = model(data_train.x_dict, data_train.edge_index_dict)  # Perform a single forward pass.      
  File "C:\Users\miniconda3\lib\site-packages\torch\fx\graph_module.py", line 616, in wrapped_call
    raise e.with_traceback(None)
AssertionError

I am genuinely confused as to why the wrapped_call(self, *args, **kwargs) would throw an Exception, because it does not really give a hint, on where the mistake could be...

Help would be greatly appreciated!

Best regards, Zytrus

[rusty1s]

Thanks for reporting and sorry for this mysterious error message. My assumption is that the metadata does not match across all your graphs (which is a requirement). That means, the following should pass:

for data in dataset:
    assert data.metadata() == dataset[0].metadata()

[rusty1s]

This is hard to tell. Can you tell me more about how you apply loss function and evaluation metrics?

[Zytrus]

I used the training from the example given in the documentation, although adjusted it a little:

criterion = torch.nn.CrossEntropyLoss()
def train():
    model.train()

    total_examples = total_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        # batch = batch.to('cuda:0')
        batch_size = batch.num_graphs
        out = model(batch.x_dict, batch.edge_index_dict, batch.batch_dict)
        loss = criterion(out[:batch_size], batch.y[:batch_size])
        loss.backward()
        optimizer.step()

        total_examples += batch_size
        total_loss += float(loss) * batch_size

    return total_loss / total_examples

def test(loader):
    model.eval()

    correct = 0
    for data_test in loader:  # Iterate in batches over the training/test dataset.
        out = model(data_test.x_dict, data_test.edge_index_dict, data_test.batch_dict)  
        pred = out.argmax(dim=1)  # Use the class with highest probability.
        correct += int((pred == data_test.y).sum())  # Check against ground-truth labels.
    return correct / len(loader.dataset)

[rusty1s]

I assume that the slicing in out[:batch_size] and batch.y[:batch_size] is not really needed. Otherwise, this looks good to me.

[Zytrus]

Thank you for the feedback! I suppose I'll just have to find a more fitting model with different layers then for better results?

Also you've been a great help with my progress, thank you very much, it is very appreciated! (:

[rusty1s]

Yes, I suggest to do some design space exploration (#layers, #hidden size, LR, different global pooling operators, ..) to see if you can improve results.