How to create Custom Datasets from saved pickle file

[Byun-jinyoung]

Dear pyg users,
Recently I made my graph data and converted it into a pickle file, which is 14G in size.
When I tried to load this data before training model, I encountered a memory-out error.
As long as I understand the InMemoryDataset needs edge, graph_idx, graph_label, node_attr, separately.
But I want to make my data in a file.
How can I solve this problem?

[rusty1s]

Can you give some more information about what you are doing concretely? Does 14GB denote the size of your data object? IMO, you should well be able to load and save that data without PyTorch or PyG needing to replicate it.

[Byun-jinyoung]

Thank you for reply! @rusty1s
I've tried to protein-ligand graph. a Protein-ligand graph has a shape like this:
Data(x=[18974, 40], y=[1], pos=[18974, 3], edge_index=[2, 711312], edge_attr=[711312, 7])
I converted the 1320 protein-ligand complex into 1320 graphs and save the 1320 graphs as one pickle.dump file. The size of the pickle-saved file is 14GB.

[rusty1s]

Ok. If you save the data inside pickle as torch.Tensor, their data should be shared when creating PyG data objects. There is an additional overhead though when calling self.collate (which groups all graphs into one big graph), which might explain the OOM error. Currently, there exists no way around this except not making use of InMemoryDataset.

[Byun-jinyoung]

Could you recommend which is the best way efficiently to save pyg Data? not to use pickle? How about saving a graph, respectively with torch.save?

[rusty1s]

One option is to utilize torch_geometric.data.Dataset and save each graph individually to disk (see here) and apply a large number of workers in DataLoader to not make I/O become a bottleneck.

[Byun-jinyoung]

Thank you for the advice!