How to build custom in memory dataset with multiple files

[Ripper346]

Hi, I am trying to make a custom Dataset class. I followed the documentation for the in memory dataset but now I need to replace the last row of init of the class to adapt it for multiple loads because I have several files to load in the same dataset. There isn't a file per graph, instead I can have 1000 graphs in data_0.pt and 300 in data_1.pt.
From what I understood (I can’t find in the docs), torch.load(file_path) returns a tuple with Data and slice. How can I combine these two for multiple files keeping the object structure? It is a little hard to explain it, I hope you understand my problem. Thank you.

I tried to combine buffers with

buffer = io.BytesIO()
for file in self.processed_paths:
    with open(file, 'rb') as f:
        buffer.write(f.read())
buffer.seek(0)
self.data, self.slices = torch.load(buffer)

but it gives the error RuntimeError: [enforce fail at ..\caffe2\serialize\inline_container.cc:145] . PytorchStreamReader failed reading zip archive: invalid header or archive is corrupted that well, is understandable due the structure of the buffer.

Then I came up using Batch class

datas = []
for file in self.processed_paths:
    data, slices = torch.load(file)
    datas.append(data)
self.data = Batch.from_data_list(datas)   # or at least a starting point

but it gives an error on row 122 of batch.py (after starting from_data_list method) TypeError: full(): argument 'size' must be tuple of ints, but found element of type list at pos 1.

Edit: I am trying also without in memory dataset approach and it seems that the files are assumed to be one graph - one file. I am completely lost in this.

[rusty1s]

In general, all you have to care about is how to create a list of data objects in process, and PyG will take care of collating, saving and loading the processed data.

I'm not exactly sure what you are trying to do, but I suggest to put all graphs of all files in a single data list, and save/load it via

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

In case you really want to save them into distinct files, you could do the following. In process, do something like:

for filename in ...:
    data_list = []  # Create a list of data objects for the file `filename`
    torch.save(self.collate(data_list), osp.join(self.processed_dir, filename))

and load one of them in __init__:

self.data, self.slices = torch.load(osp.join(self.processed_dir, filename))

[Ripper346]

Yes, I did it as you say for saving the processed data in different files. I am in the last piece of code you wrote, why only one? Can't I load into self.data, self.slices all the files I saved in the process?

[rusty1s]

No, such a behavior is currently not supported for our provided dataset interface, but you could easily set that up by yourself by inheriting from torch.utils.data.Dataset directly.

[Ripper346]

So I can't inherit my class from InMemoryDataset or torch_geometric.data.Dataset. From torch I don't know how to do it, I have to look at it

[Ripper346]

Self answering, I almost forgot, maybe there are others with the same problem. I came up looking for too long to the examples of in memory and not in memory classes. The trick to load multiple file is the place of the collate function call. In the example there is

import torch
from torch_geometric.data import InMemoryDataset, download_url


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 process(self):
        # Read data into huge `Data` list.
        data_list = [...]

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

and

import os.path as osp

import torch
from torch_geometric.data import Dataset, download_url


class MyOwnDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyOwnDataset, self).__init__(root, transform, pre_transform)

    def process(self):
        i = 0
        for raw_path in self.raw_paths:
            # Read data from `raw_path`.
            data = Data(...)

            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data

In the second there is something missing when saving in process and surprise surprise it's the collate as I wrote in the beginning. So how to load multiple files in memory? Like this:

import torch
from torch_geometric.data import InMemoryDataset, download_url


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

    def process(self):
        # Read data into huge `Data` list.
        data_list = [...]
        torch.save(data_list, self.processed_paths[0])

This example is the same but if we save more files in process we can concatenate data_list in the constructor and do only one collate in the end. Like:

class MyOwnDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyOwnDataset, self).__init__(root, transform, pre_transform)
        data_list = []
        for path in self.processed_paths:
            data_list += torch.load(path)
        self.data, self.slices = self.collate(data_list)
   # etc