Repository: pyg-team/pytorch_geometric
Pull request #7250 (2023-04-27 17:21:11)

Introduce `InMemoryDataset.save` and `InMemoryDataset.load`

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7250](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (3a51301) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/5778c65686b0c6af40e745a3ec449e7b27628ead?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (5778c65) will **decrease** coverage by `0.43%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head 3a51301 differs from pull request most recent head 720cfd6. Consider uploading reports for the commit 720cfd6 to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7250      +/-   ##
==========================================
- Coverage   92.04%   91.62%   -0.43%     
==========================================
  Files         437      437              
  Lines       24145    24167      +22     
==========================================
- Hits        22225    22143      -82     
- Misses       1920     2024     +104     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/data/hetero\_data.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvaGV0ZXJvX2RhdGEucHk=) | `91.17% <100.00%> (+0.20%)` | :arrow_up: |
| [torch\_geometric/data/in\_memory\_dataset.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvaW5fbWVtb3J5X2RhdGFzZXQucHk=) | `97.97% <100.00%> (+0.25%)` | :arrow_up: |
| [torch\_geometric/data/storage.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvc3RvcmFnZS5weQ==) | `83.02% <100.00%> (-0.35%)` | :arrow_down: |

... and [19 files with indirect coverage changes](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7250/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 163647b09..abc96f648 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `save` and `load` methods to `InMemoryDataset` ([#7250](https://github.com/pyg-team/pytorch_geometric/pull/7250))
 - Added an example for heterogeneous GNN explanation via `CaptumExplainer` ([#7096](https://github.com/pyg-team/pytorch_geometric/pull/7096))
 - Added `visualize_feature_importance` functionality to `HeteroExplanation` ([#7096](https://github.com/pyg-team/pytorch_geometric/pull/7096))
 - Added a `AddRemainingSelfLoops` transform ([#7192](https://github.com/pyg-team/pytorch_geometric/pull/7192))
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 7a9144e51..abca8fd99 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -125,6 +125,23 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
             else:
                 setattr(self, key, value)
 
+    @classmethod
+    def from_dict(cls, mapping: Dict[str, Any]) -> 'HeteroData':
+        r"""Creates a :class:`~torch_geometric.data.HeteroData` object from a
+        Python dictionary."""
+        out = cls()
+        for key, value in mapping.items():
+            if key == '_global_store':
+                out.__dict__['_global_store'] = BaseStorage(
+                    _parent=out, **value)
+            elif isinstance(key, str):
+                out._node_store_dict[key] = NodeStorage(
+                    _parent=out, _key=key, **value)
+            else:
+                out._edge_store_dict[key] = EdgeStorage(
+                    _parent=out, _key=key, **value)
+        return out
+
     def __getattr__(self, key: str) -> Any:
         # `data.*_dict` => Link to node and edge stores.
         # `data.*` => Link to the `_global_store`.
@@ -258,11 +275,12 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         return list(self._edge_store_dict.items())
 
     def to_dict(self) -> Dict[str, Any]:
-        out = self._global_store.to_dict()
+        out_dict: Dict[str, Any] = {}
+        out_dict['_global_store'] = self._global_store.to_dict()
         for key, store in chain(self._node_store_dict.items(),
                                 self._edge_store_dict.items()):
-            out[key] = store.to_dict()
-        return out
+            out_dict[key] = store.to_dict()
+        return out_dict
 
     def to_namedtuple(self) -> NamedTuple:
         field_names = list(self._global_store.keys())
diff --git a/torch_geometric/data/in_memory_dataset.py b/torch_geometric/data/in_memory_dataset.py
index 1e1b3ee19..9d54cf1b7 100644
--- a/torch_geometric/data/in_memory_dataset.py
+++ b/torch_geometric/data/in_memory_dataset.py
@@ -2,12 +2,24 @@ import copy
 import warnings
 from abc import ABC
 from collections.abc import Mapping, Sequence
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
-
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
+
+import torch
 from torch import Tensor
 
 from torch_geometric.data import Batch, Data
 from torch_geometric.data.collate import collate
+from torch_geometric.data.data import BaseData
 from torch_geometric.data.dataset import Dataset, IndexType
 from torch_geometric.data.separate import separate
 
@@ -93,6 +105,19 @@ class InMemoryDataset(Dataset, ABC):
 
         return data
 
+    @classmethod
+    def save(cls, data_list: List[BaseData], path: str):
+        r"""Saves a list of data objects to the file path :obj:`path`."""
+        data, slices = cls.collate(data_list)
+        torch.save((data.to_dict(), slices), path)
+
+    def load(self, path: str, data_cls: Type[BaseData] = Data):
+        r"""Loads the dataset from the file path :obj:`path`."""
+        data, self.slices = torch.load(path)
+        if isinstance(data, dict):  # Backward compatibility.
+            data = data_cls.from_dict(data)
+        self.data = data
+
     @staticmethod
     def collate(
             data_list: List[Data]) -> Tuple[Data, Optional[Dict[str, Tensor]]]:
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index 41a68c699..f43cf1254 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -191,7 +191,13 @@ class BaseStorage(MutableMapping):
 
     def to_dict(self) -> Dict[str, Any]:
         r"""Returns a dictionary of stored key/value pairs."""
-        return copy.copy(self._mapping)
+        out_dict = copy.copy(self._mapping)
+        # Needed to preserve individual `num_nodes` attributes when calling
+        # `BaseData.collate`.
+        # TODO (matthias) Try to make this more generic.
+        if '_num_nodes' in self.__dict__:
+            out_dict['_num_nodes'] = self.__dict__['_num_nodes']
+        return out_dict
 
     def to_namedtuple(self) -> NamedTuple:
         r"""Returns a :obj:`NamedTuple` of stored key/value pairs."""


Test patch:
diff --git a/test/data/test_dataset.py b/test/data/test_dataset.py
index 067a055a1..d82cb09bd 100644
--- a/test/data/test_dataset.py
+++ b/test/data/test_dataset.py
@@ -10,10 +10,24 @@ from torch_geometric.typing import SparseTensor
 
 class MyTestDataset(InMemoryDataset):
     def __init__(self, data_list, transform=None):
-        super().__init__('/tmp/MyTestDataset', transform=transform)
+        super().__init__(None, transform=transform)
         self.data, self.slices = self.collate(data_list)
 
 
+class MyStoredTestDataset(InMemoryDataset):
+    def __init__(self, root, data_list, transform=None):
+        self.data_list = data_list
+        super().__init__(root, transform=transform)
+        self.load(self.processed_paths[0], data_cls=data_list[0].__class__)
+
+    @property
+    def processed_file_names(self) -> str:
+        return 'data.pt'
+
+    def process(self):
+        self.save(self.data_list, self.processed_paths[0])
+
+
 def test_in_memory_dataset():
     x1 = torch.Tensor([[1], [1], [1]])
     x2 = torch.Tensor([[2], [2], [2]])
@@ -57,6 +71,54 @@ def test_in_memory_dataset():
     assert torch.equal(dataset[1:].x, x2)
 
 
+def test_stored_in_memory_dataset(tmp_path):
+    x1 = torch.Tensor([[1], [1], [1]])
+    x2 = torch.Tensor([[2], [2], [2], [2]])
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    data1 = Data(x1, edge_index, num_nodes=3, test_int=1, test_str='1')
+    data2 = Data(x2, edge_index, num_nodes=4, test_int=2, test_str='2')
+
+    dataset = MyStoredTestDataset(tmp_path, [data1, data2])
+    assert dataset._data.num_nodes == 7
+    assert dataset._data._num_nodes == [3, 4]
+
+    assert torch.equal(dataset[0].x, x1)
+    assert torch.equal(dataset[0].edge_index, edge_index)
+    assert dataset[0].num_nodes == 3
+    assert torch.equal(dataset[0].test_int, torch.tensor([1]))
+    assert dataset[0].test_str == '1'
+
+    assert torch.equal(dataset[1].x, x2)
+    assert torch.equal(dataset[1].edge_index, edge_index)
+    assert dataset[1].num_nodes == 4
+    assert torch.equal(dataset[1].test_int, torch.tensor([2]))
+    assert dataset[1].test_str == '2'
+
+
+def test_stored_hetero_in_memory_dataset(tmp_path):
+    x1 = torch.Tensor([[1], [1], [1]])
+    x2 = torch.Tensor([[2], [2], [2], [2]])
+
+    data1 = HeteroData()
+    data1['paper'].x = x1
+    data1['paper'].num_nodes = 3
+
+    data2 = HeteroData()
+    data2['paper'].x = x2
+    data2['paper'].num_nodes = 4
+
+    dataset = MyStoredTestDataset(tmp_path, [data1, data2])
+    assert dataset._data['paper'].num_nodes == 7
+    assert dataset._data['paper']._num_nodes == [3, 4]
+
+    assert torch.equal(dataset[0]['paper'].x, x1)
+    assert dataset[0]['paper'].num_nodes == 3
+
+    assert torch.equal(dataset[1]['paper'].x, x2)
+    assert dataset[1]['paper'].num_nodes == 4
+
+
 def test_in_memory_num_classes():
     dataset = MyTestDataset([Data(), Data()])
     assert dataset.num_classes == 0
diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py
index 1ea3c79bb..e0be1dd69 100644
--- a/test/data/test_hetero_data.py
+++ b/test/data/test_hetero_data.py
@@ -90,6 +90,20 @@ def test_init_hetero_data():
     assert len(data.edge_items()) == 3
 
 
+def test_hetero_data_to_from_dict():
+    data = HeteroData()
+    data.global_id = '1'
+    data['v1'].x = torch.randn(5, 16)
+    data['v2'].y = torch.randn(4, 16)
+    data['v1', 'v2'].edge_index = torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]])
+
+    out = HeteroData.from_dict(data.to_dict())
+    assert out.global_id == data.global_id
+    assert torch.equal(out['v1'].x, data['v1'].x)
+    assert torch.equal(out['v2'].y, data['v2'].y)
+    assert torch.equal(out['v1', 'v2'].edge_index, data['v1', 'v2'].edge_index)
+
+
 def test_hetero_data_functions():
     data = HeteroData()
     data['paper'].x = x_paper