Break batch according to number of nodes and then recombine it as it was

[antoni4040]

A bit of a head-scratching riddle for an experiment I want to try. I have a batch with let's say 100 graphs. Some are large, some are medium, some are small, depending on the number of nodes. I want to break the batch, pass each of these though different GNNs (according to their size category) and then recombine the results. I could of course just make batches of one, but I'm worried about the slowdown. Thank you and sorry if it's a bit weird :)

[rusty1s]

If you want to apply different GNNs depending of size, how about you separate these data objects in a pre-processing step and train your models sequentially? If you want to do it within one model, I think you have two options:

1. Define a custom sampler on the DataLoader that only samples graphs of equal size. This way, each forward method will only need to use a single GNN.
2. Separate your graphs during forward execution, e.g., via data_list = batch.to_data_list(), split them up, and re-combine to separate batches via Batch.from_data_list(). Note that this will likely be slow.