"Uncollate" a batch

[anticdimi]

Hi,

I was wondering whether it is possible to somehow get a batch in "pytorch style" out of the batch that is yielded by torch_geometric.loader.DataLoader? ie. if I have X point clouds, each consisting of eg. Y 3D points, I end up with a batch of size (X*Y, 3). What would be the easiest way to obtain a corresponding batch of size (X, Y, 3)?

Thanks!
Dimitrije

[rusty1s]

If your points are all of the same size, simply reshaping is sufficient:

x = x.view(batch_size, num_nodes, 3)

Otherwise, take a look at utils.to_dense_batch, which will do the same with some padding logic.

[anticdimi]

Hi @rusty1s,

That seemed like a solution to me too, but I wondered if the order of points/nodes is preserved.
Thanks!

Best,
Dimitrije

[rusty1s]

Yes, the ordering is preserved.