GCNConv,pass a Tensor of size [GraphNum, PointNum, FeatureDim] and [2, EdgeNum] in the forward method

[wishForget]

Hello everyone!
I'm a beginner of PyG and I have search for my problems but found nothing useful. I encountered a situation while using GCNConv where I attempted to pass a Tensor of size [GraphNum, PointNum, FeatureDim] in the forward method, representing batches of graphs with GraphNum nodes and PointNum points each (as opposed to the [PointNum, FeatureDim] size tensor as specified in the documentation). Additionally, I passed an edge tensor of size [2, EdgeNum]. Surprisingly, the program ran successfully and the method has worked.
However, when I tried the same approach with GATConv, the program did not run successfully.
I would like to understand what is happening in this scenario. Thank you everyone!!!!!!!!
Here is my code.

class GraphNeuralNetwork(nn.Module):

    def __init__(
            self,
            num_layers: int,
            input_dim: int,
            hidden_dim: int,
    ):
        super(GraphNeuralNetwork, self).__init__()
        self.num_layers = num_layers
        self.gcn_layers = nn.ModuleList()
        self.gcn_layers.append(gnn.GraphConv(input_dim, hidden_dim))
        self.gcn_layers.append(nn.ReLU(inplace=True))
        self.gcn_layers.append(gnn.GraphNorm(hidden_dim))
        for _ in range(num_layers):
            self.gcn_layers.append(gnn.GraphConv(hidden_dim, hidden_dim))
            self.gcn_layers.append(nn.ReLU(inplace=True))
            self.gcn_layers.append(gnn.GraphNorm(hidden_dim))

        self.regression_head = nn.Sequential(
                nn.Linear(hidden_dim, 64),
                nn.ReLU(inplace=True),
                nn.Linear(64, 2),
            )

    def forward(
            self,
            x: torch.Tensor,  #size  [graph_num, point_num, dim]
            edge_index: torch.Tensor  #size [2, edge_num]
    ) -> Dict[str, torch.Tensor]:
        for layer in self.gcn_layers:
            if isinstance(layer, gnn.GraphConv):
                x = layer(x, edge_index)
            else:
                x = layer(x)

def test():
    num_layers = 1
    input_dim = 64
    hidden_dim = 128
    B, N, C = 8, 64, input_dim
    x = torch.randn(B, N, C) 
    edges = torch.tensor([(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7), (7, 8), (8, 9)], dtype=torch.long).permute(1,0)
    gnn_model = GraphNeuralNetwork(num_layers, input_dim, hidden_dim)
    res = gnn_model(x, edge_index)

if __name__ == '__main__':
    test()

[rusty1s]

There exists a few layers in PyG that support this "static graph" computation where the graph is static across a set of batches, and thus allow for [batch_size, num_nodes, num_features] feature tensors. Supported layers are documented here in the "static" column.

[rusty1s]

Hey, yes, the recommended way is as you describe, e.g.:

data1 = Data(xyz=torch.randn(x, y, z), ...)
data2 = Data(xyz=torch.randn(x, y, z), ...)

loader = DataLoader([data1, data2, ...], batch_size=...)

for batch in loader:
    print(batch.xyz.shape)

This should yield batch.xyz with shape [2 *x, y, z].

[bet20ICL]

Thanks so much for clarifying! I'm still slightly unclear on some points:

In my case, each training sample is accesssed from a Dataset object. The __getitem__ method of the Dataset loads each training sample from a file.

When I wrap my Dataset using the PyG DataLoader, each batch comes out with dim (BATCH, X, Y, Z) rather than (BATCH*X, Y, Z). Is this expected behaviour?

[rusty1s]

No, it‘s not. Do you mind sharing a small example?

[bet20ICL]

Of course! Thanks so much for looking into this. This is an abridged version of the code I'm using:

class WeatherDataset(torch.utils.data.Dataset):
    ...
    def __getitem__(self, idx):
        # === Sample ===
        sample_path = os.path.join(self.sample_dir_path, f"nwp_{idx}.npy")
        # load the following variables from the file at sample_path
        ...
        # for example let init_states.shape = (x, y, z)
        return init_states, target_states, static_features, forcing_windowed

args.batch_size = 2
train_loader = pyg.loader.DataLoader(
            WeatherDataset(
                args.dataset, ...),
            args.batch_size, ...)

for batch in train_loader:
    batch[0].shape 
    # Expected: [2*x, y, z]
    # Actual: [2, x, y, z]

[rusty1s]

Thanks. Note that you don't use a Data object here, so the default collating of mini-batches is picked up by PyTorch. Try

def __getitem__(self, idx):
    return Data(init_states=init_states, ...)