pytorch-geometricmini-batchgraph-neural-network

How to make a minibatch with different graphs in GNN


How to make a minibatch with different graphs in GNN so I can aggregate their information respectively in their own graph.

If I have four clips of x([1024,1]), and each clip of x has its own graph(every graph is different), how can I add them together as a minibatch so as to aggregate their information respectively with GNN in their own graph.

For example, if we define 2 sets of x:

    x_a = torch.randn(2, 16)  # 2 nodes.
    x_b = torch.randn(3, 16)  # 3 nodes.

    edge_index_a = torch.tensor([[0, 0],
                               [0, 1]])

    edge_index_b = torch.tensor([[0, 0, 1, 1],
                               [0, 1, 1, 2]])

How can I add them together as PairData(x[5,16], edge_index[2,6]) and the shape of added edge_index like

([0, 0, 2, 2, 3, 3],                                        
 [0, 1, 2, 3, 3, 4]) 

Solution

  • The DataLoader from torch_geometric.loader do this for you. You don't need to implement it by yourself. Check this example on graph classification on molecular property prediction