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)
edge_index
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)
print(x.size())
torch.Size([7, 16])
and pass it through gconv
, I have:
y = gconv(x, edge_index)
print(y.size())
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.