Using custom collate_fn with PyG DataLoader

[hatemhelal]

👋 I wonder if anyone can share how to use a custom collate_fn with the PyG DataLoader. From the code it appears that we need to provide our own implementation since the collate_fn is removed here

I'd be interested in hearing of any other solutions.

Thanks!

[wsad1]

That line you are referring to removes user defined collate functions, the dataloader instead uses a inbuilt collate function to collate graphs, refer to this line.

Now you could modify the "behaviour" of the default collate function by over-ridding the __inc__ and __cat_dim__ functions as described (here)[https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html#pairs-of-graphs]. Now if that too doesn't fit your use case you'll have to write a custom dataloader with your own collate function.

[rusty1s]

Yes, you can simply use torch.utils.data.DataLoader to implement your own collate_fn (this is exactly what we are doing within our PyG DataLoader).

[hatemhelal]

Thanks both for the helpful answers. I think for my use case the __inc__ and __cat_dim__ aren't sufficient since I am injecting padding nodes and edges to ensure consistent tensor sizes across an entire dataset.

I had the idea of adding an optional collate_fn argument to the PyG DataLoader would would default to the current Collater. Let me know if this approach sounds worth pursuing and I can try to make a PR. 🙏

[wsad1]

PyG can handle batching of graphs with different number of nodes and edges, as explained here.
But feel free to elaborate on your problem, if the above doesn't fit your use case.