[PyTorch Geometric] Graph convolution in batch mode

Let’s say that we need to define a batch of graphs, of the same structure, i.e., with the same edge_index, but with different feature signals and different edge attributes.

For instance, let’s define a simple directed graph structure with the following edge_index:

import torch
from torch_geometric.data import Data as gData
import torch_geometric.nn as gnn
import numpy as np

edge_index = torch.tensor(np.concatenate([np.arange(num_nodes), np.roll(np.arange(num_nodes), shift=1)]).reshape(-1, num_nodes),  dtype=torch.long)
tensor([[0, 1, 2, 3, 4, 5, 6],
        [6, 0, 1, 2, 3, 4, 5]])

Now, let’s define a simple graph convolution operator, e.g., GCNConv, that will act on such graphs:

gconv = gnn.GCNConv(in_channels=num_node_features, out_channels=32)

Then, if I define a graph signal as below:

x = torch.randn((num_nodes, num_node_features), dtype=torch.float)
torch.Size([7, 16])

and pass it through gconv, I have:

y = gconv(x, edge_index)
torch.Size([7, 32])

which is fine.

Now, I’d like to do the same in a mini-batch manner; i.e., to define a a batch of such signals, that along with the same edge_index will be passed through gconv. Apparently, defining signals and edge attributes as 3D tensors does not work.

I suspect that this could be somehow done using batch, mentioned here, but I cannot find any reference on how this could be done.

1 Like

I have the same problem. Have you found any solutions?

You can refer to this issue in the original repository.

1 Like