I am creating a message passing neural network and have some issues with the dataset creation. In pytorch (geometric) it is recommended to create a dataset with the following class. I wonder what is the meaning of the collate function that is called at the end of the process method? In what cases should I use my own collate function? My graphs have mostly different sizes.

import torch
from torch_geometric.data import InMemoryDataset

class MyOwnDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyOwnDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    def raw_file_names(self):
        return ['some_file_1', 'some_file_2', ...]

    def processed_file_names(self):
        return ['data.pt']

    def download(self):
        # Download to `self.raw_dir`.

    def process(self):
        # Read data into huge `Data` list.
        data_list = [...]

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

The error I get is:

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 4422 and 4032 in dimension 1 at /opt/conda/conda-bld/pytorch_1573049304260/work/aten/src/THC/generic/THCTensorMath.cu:71


1 Answer 1


What collate does and why:

Because saving a huge python list is really slow, we collate the list into one huge torch_geometric.data.Data object via torch_geometric.data.InMemoryDataset.collate() before saving . The collated data object has concatenated all examples into one big data object and, in addition, returns a slices dictionary to reconstruct single examples from this object. Finally, we need to load these two objects in the constructor into the properties self.data and self.slices.


  • $\begingroup$ I am nto a native speaker and not familiar with the concept of 'collating'. The dictionary sais "to bring together different pieces of written information so that the similarities and differences can be seen". That seems to be something else in this case. Maybe it is better described as concatenating? $\endgroup$
    – Soerendip
    Jan 15, 2021 at 2:27

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.