collate function for handling non tensor data

[ipcamit]

Hi,
In Pytorch I can give collate function to convert my data (list of molecule xyz data structures) to desired torch tensor using the collate_fn argument in DataLoader. Is there some way to do it in torch_geometric.loader.DataLoader directly?
i.e. currently I do:

dl = torch.data.DataLoader(data_set, batch_size=2, collate_fn=custom_collate_fn)
pydl = next(iter(torch_geometric.loader.DataLoader(next(iter(dl)), batch_size=2))

to get properly collated joint graph. Here custom_collate_fn takes in list of my custom data class, and returns list of individual torch_geometric.data.Data classes. Is there any way torch_geometric.loader.DataLoader can take this collate function instead?

I tried

pydl = torch_geometric.loader.DataLoader(data_set, batch_size=2,collate_fn=custom_collate_fn )

in hope that it will get passed to underlying torch dataloader. But I got following error:

File /opt/anaconda3/anaconda3/envs/colabfit/lib/python3.9/site-packages/torch_geometric/loader/dataloader.py:36, in Collater.__call__(self, batch)
     33 elif isinstance(elem, Sequence) and not isinstance(elem, str):
     34     return [self(s) for s in zip(*batch)]
---> 36 raise TypeError(f'DataLoader found invalid type: {type(elem)}')

TypeError: DataLoader found invalid type: <class 'kliff_torch.dataset.dataset.Configuration'>

[EdisonLeeeee]

collate_fn is not supported for torch_geometric.loader.DataLoader

pytorch_geometric/torch_geometric/loader/dataloader.py

Lines 71 to 72 in d5e65e6

	if 'collate_fn' in kwargs:
	del kwargs['collate_fn']

[rusty1s]

Indeed. Our DataLoader works by overriding collate_fn. If you want to define collate_fn on your own, simply use torch.utils.data.DataLoader.

[ipcamit]

I understood that, I was wondering if there is any procedure i can call before torch_geometric collate_fn is called. So that I can apply some simple transform to data first.

[rusty1s]

Interesting. This is currently handled outside the data loader as part of transforms in the Dataset class. Does this work for your use-case?

[ipcamit]

Thanks for the suggestion. It worked as I expected. Below is very barebones Dataset I wrote which worked as expected:

class SimplePYGDataset(PyGDataset):
    def __init__(self, data, transform):
        super().__init__("./", transform, None, None)
        self.data_list = data
    
    def len(self):
        return len(self.data_list)
    
    def get(self, idx):
        return [self.data_list[idx]]

Is there a simple way of exporting/saving the complete dataset? Currently I create InMemoryDataset then save it using torch.save(imd.collate(imd.data)).

[rusty1s]

You can also just save the InMemoryDataset, e.g., via torch.save(dataset).