GraphLayers + LinearLayers combination

[AlejandroTL]

Hello!

I'm working on an architecture that combines graph layers and linear layers. For the graph layers with message passing, the way in which the mini-batches are built is fine. But when reaching the linear layers, I want to concatenate the node features of each graph and input that node features concatenated to the linear layers. The problem is that I don't find a way to extract just the features of a single graph.

The mini-batch is a single graph with different connected components that are the graphs itself. How can I isolate just one graph each time to input it to the nn.Linear layers at each time? Easiest solution would be use ' batch_size=1` but I don't want to do it. I've looking for some info or hints in this discussion forum but found nothing :(

My pseudocode is something like this:

GCNConv(in,out)
GCNConv(out,out/2)
-----> here I want to concatenate the node features of each graph
nn.Linear(x,z)

The ideal solution could be something like a pooling_layer but to concatenate, instead of perform the add, mean or max.
Any idea?

Thank you!

[rusty1s]

I think this is achievable via to_dense_batch, which returns a tensor of shape [batch_size, max_num_nodes, num_features] which you can reshape to [batch_size, max_num_nodes * num_features]. Keep in mind that it will pad with zero features in case your graphs are of varying size.