Batched DGLGraphs with GraphConv

I’ve been trying to understand the approach for minibatch training with a DGLGraph. It seems the recommended approach is to use dgl.batch(<graph_list>) which then creates a completely new graph and combines each graphs ndata and edata.

What is confusing here, is that looking at the implementation for GraphConv, I’m to provide a graph and node features. If I do this with a batched graph I would expect to provide the node features with an additional batch dimension, but this is not supported.

It seems when working with batched graphs, we should have all the features in the graph before batching, which is a completely different idiom from what’s recommended for non-batched graphs. The problem with doing this, is it seems the only way to get results back from a batched graph is to use the provided readout functions. In my case, I want a per node value with dimensions (batch, node, features). The only way I can see to do this is to return the graph, call dgl.unbatch(graph) and then extract the node values. Not only is this cumbersome, but I’m not sure if it causes any problems for the autograd engine (using Pytorch)

TLDR: Want to be able to provide batch graphs and batched features to a GraphConv and get batched features back for training.

Additional Point: I can get features back from BatchedGraph but I’m not sure how to reshape the data from (B*N, F) to (B, N, F) since I don’t know for sure the order that batching is done.

2 Likes

No, this is not what will happen, you only need to provide node features of shape (N1+N2+…N_b, d) where b refers to the batch size.

You may refer to this doc to understand the layout of node/edge features in batched DGLGraph.

https://www.dgl.ai/blog/2019/01/25/batch.html is an excellent post showing how to use GraphConv and readout modules on batched graphs.

@pbontrager Do all your graphs have the same number of nodes? The reason why dgl.batch returns a new graph with combined ndata and edata is because graphs can have different nodes and edges. Thus, I cannot see how to represent each node value with a tensor of shape (batch, node, features), because the length of the second dimension is not fixed.

If all your graphs happen to have the same number of nodes, it actually becomes much easier to get results back from a batched graph by reshaping.

all_node_feats = ...  # tensor of shape (batch, num_nodes, features)
G = dgl.batch(graphs)  # the batched graph has num_nodes*batch nodes
batched_feats = all_nodes_feats.view(-1, all_nodes_feats.shape[2])  # shape (batch*num_nodes, features)
rst = conv(G, batched_feats)   # apply GraphConv on the batched graph
new_all_node_feats = rst.view(batch, num_nodes, features)  # no unbatch is needed

The new graph returned by dgl.batch has each graph as an isolated component. As a result, feeding this batched graph to GraphConv automatically performs graph convolution on multiple graphs in one shot.

2 Likes

My graphs do all happen to have the same exact structure, so this answer is perfect, thank you. I just wasn’t sure at first if the reshaping would match exactly how dgl.batch did the reshape.

You may refer to this doc to understand the layout of node/edge features in batched DGLGraph.
https://www.dgl.ai/blog/2019/01/25/batch.html is an excellent post showing how to use GraphConv and readout modules on batched graphs.

This post is what I originally followed but it doesn’t account for node and edge features when you aren’t using a readout module. In this case it would be very helpful to have a readout module that just returns your node or edge features. In the case of a batched graph it could return node or edge features with a batch dimension and automatically pad the node/edge dimension so every graph output would be the same size. Thank you for getting back to me though.