The batch format for GATConv

[Fujiry0]

Hello. I have a question about the batch format in GATConv. I'm writing the code of GAT for a fully connected garaph with reference to some issues (e.g. here and here), but it seems this model does not learn anything. I think it is caused by my misunderstanding about the way to make batches. Is there are any best practices to deal with the batch?

import torch
from torch import nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv, GraphConv, GCNConv, AGNNConv, EdgeConv, DenseSAGEConv
from torch_geometric.data import Data as gData
from torch_geometric.data import Batch
import time
cuda = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class GAT(nn.Module):
    def __init__(self, in_channels=1024, out_channels=1024):
        super().__init__()
        self.conv1 = GATConv(in_channels,  out_channels)
        self.conv2 = GATConv(out_channels, out_channels)

    def forward(self, x):
        batch_size, num_nodes, fc_dim = x.shape
        edge_index_list = self.build_edge_index(batch_size, num_nodes)
        graph = self.build_graph_batch(x, edge_index_list)
        x = self.conv1(graph.x, graph.edge_index)
        x = F.relu(x)
        x = self.conv2(x, graph.edge_index)
        return x

    def build_edge_index(self,  batch_size, num_nodes):
        row = torch.arange(num_nodes).view(-1, 1).repeat(1, num_nodes).view(-1)
        col = torch.arange(num_nodes).view(-1, 1).repeat(num_nodes, 1).view(-1)
        edge_index = torch.stack([row, col], dim=0).to(cuda)
        edge_index_list = [edge_index for _ in range(batch_size)]
        return edge_index_list

    def build_graph_batch(self, x_list,  edge_index_list):
        data_list = []
        for graph_nodes, ei in zip(x_list, edge_index_list):
            data_list.append(gData(x=graph_nodes, edge_index=ei))
        graph = Batch.from_data_list(data_list)
        return graph


if __name__ == '__main__':
    cuda = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    batch_size = 2
    num_nodes = 128
    fc_dim = 512
    model = GAT(fc_dim, fc_dim).to(cuda)
    x = torch.randn((batch_size, num_nodes, fc_dim), device=cuda)
    out = model(x)
    print(out.shape)

Thanks in advance.

[rusty1s]

I think the way you create batches is correct (although for best-practice you should separate it from the model instance). Therefore, I think the problem of not learning anything needs to stem from something else.

[Fujiry0]

Thank you so much for your help! I will check other components.

Thanks for the great library.

[wangsir2]

Do I need to add F.relu between self.conv1 and self.conv1?

[rusty1s]

Yes (or any other non-linearity). We do not apply this inside the GNN layers.

[wangsir2]

Thanks for Your Help