Using Dataset class still gets "Cannot allocate memory" error

[LanceKnight]

Hello,

I'm not so familiar how memory is allocated. But I thought the purpose of Dataset class is to let a big dataset fit into the memory, as compared with InMemoryDataset that cannot do this. However, I still got "cannot allocate memory" for big dataset. Can someone remind me of what could have been done wrong? Or if my understanding of Dataset is incorrect? Thanks!

[rusty1s]

In contrast to InMemoryDataset, Dataset will load graphs individually in memory (think of a large set of graphs), such that one can process a dataset which overall does not fit into CPU memory. However, individual graphs still need to fit into memory. There is currently no elegant solution to fix this. The only solution I am aware of is to make use of numpy.memmap to store node features. Memory-mapped files are used for accessing small segments of large files on disk, without reading the entire file into memory.

[rusty1s]

So that is an in-memory dataset, right? It holds all graphs in-memory inside data_list and runs OOM once data_list gets too big.

[LanceKnight]

Oops, sorry I copied and pasted the version before I changed InMemoryDataset to Dataset. I basically change parent class to Dataset and added needed methods of len(), get() and collate() from the InMemoryDataset class:

class MyDataset(Dataset):

    def __init__(self, root, input_file):
        self.csv_file = input_file
        super(MyDataset, self).__init__(root) # Get the __init__ function in the super class, i.e., InMemoryDatset, running
        # Load the data
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def processed_file_names(self):
        return 'processed_data.pt'

    def process(self):
        data_df = pd.read_csv(self.csv_file) # Convert .csv file to pandas dataframe
        smiles_list = list(data_df['smiles']) # Convert the dataframe to a list of SMILES string.
        data_list = []

        # Process each SMILES in the .csv file
        for smiles in smiles_list:
            data = convert_SMILES_to_graph(smiles)
            data_list.append(data)

        # Save our processed data_list into processed_data.pt
        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

    def len(self):
        if self.slices is None:
            return 1
        for _, value in nested_iter(self.slices):
            return len(value) - 1
        return 0

    def get(self, idx: int) -> Data:
        if self.len() == 1:
            return copy.copy(self.data)

        if not hasattr(self, '_data_list') or self._data_list is None:
            self._data_list = self.len() * [None]
        elif self._data_list[idx] is not None:
            return copy.copy(self._data_list[idx])

        data = separate(
            cls=self.data.__class__,
            batch=self.data,
            idx=idx,
            slice_dict=self.slices,
            decrement=False,
        )

        self._data_list[idx] = copy.copy(data)

        return data

    @staticmethod
    def collate(data_list):
        r"""Collates a Python list of :obj:`torch_geometric.data.Data` objects
        to the internal storage format of
        :class:`~torch_geometric.data.InMemoryDataset`."""
        if len(data_list) == 1:
            return data_list[0], None

        data, slices, _ = collate(
            data_list[0].__class__,
            data_list=data_list,
            increment=False,
            add_batch=False,
        )

        return data, slices

[rusty1s]

Note that this just replicates the InMemoryDataset functionality. The goal is to not store any data_list. In fact, you may want to do something like:

        # Process each SMILES in the .csv file
        for i, smiles in enumerate(smiles_list):
            data = convert_SMILES_to_graph(smiles)
            torch.save(data, osp.join(self.processed_dir, f'data_{i}.pt'))

inside process().

[LanceKnight]

I see. I'm gonna give it a try. Thank you so much :)

[LanceKnight]

Worked well. Thanks!