Repository: pyg-team/pytorch_geometric
Pull request #8116 (2023-10-03 09:21:59)

`InMemoryDataset.to_on_disk_dataset()` conversion


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bdf1eae21..1f9da5e5a 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 the `to_on_disk_dataset()` method to convert `InMemoryDataset` instances to `OnDiskDataset` instances ([#8116](https://github.com/pyg-team/pytorch_geometric/pull/8116))
 - Added `torch-frame` support ([#8110](https://github.com/pyg-team/pytorch_geometric/pull/8110))
 - Added the `DistLoader` base class ([#8079](https://github.com/pyg-team/pytorch_geometric/pull/8079))
 - Added `HyperGraphData` to support hypergraphs ([#7611](https://github.com/pyg-team/pytorch_geometric/pull/7611))
diff --git a/torch_geometric/data/in_memory_dataset.py b/torch_geometric/data/in_memory_dataset.py
index 81940e4ea..f504713da 100644
--- a/torch_geometric/data/in_memory_dataset.py
+++ b/torch_geometric/data/in_memory_dataset.py
@@ -1,4 +1,5 @@
 import copy
+import os.path as osp
 import warnings
 from abc import ABC
 from collections.abc import Mapping, Sequence
@@ -16,7 +17,9 @@ from typing import (
 
 import torch
 from torch import Tensor
+from tqdm import tqdm
 
+import torch_geometric
 from torch_geometric.data import Batch, Data
 from torch_geometric.data.collate import collate
 from torch_geometric.data.data import BaseData
@@ -160,6 +163,99 @@ class InMemoryDataset(Dataset, ABC):
         dataset.data, dataset.slices = self.collate(data_list)
         return dataset
 
+    def to_on_disk_dataset(
+        self,
+        root: Optional[str] = None,
+        backend: str = 'sqlite',
+        log: bool = True,
+    ) -> 'torch_geometric.data.OnDiskDataset':
+        r"""Converts the :class:`InMemoryDataset` to a :class:`OnDiskDataset`
+        variant. Useful for distributed training and hardware instances with
+        limited amount of shared memory.
+
+        root (str, optional): Root directory where the dataset should be saved.
+            If set to :obj:`None`, will save the dataset in
+            :obj:`root/on_disk`.
+            Note that it is important to specify :obj:`root` to account for
+            different dataset splits. (optional: :obj:`None`)
+        backend (str): The :class:`Database` backend to use.
+            (default: :obj:`"sqlite"`)
+        log (bool, optional): Whether to print any console output while
+            processing the dataset. (default: :obj:`True`)
+        """
+        if root is None and (self.root is None or not osp.exists(self.root)):
+            raise ValueError(f"The root directory of "
+                             f"'{self.__class__.__name__}' is not specified. "
+                             f"Please pass in 'root' when creating on-disk "
+                             f"datasets from it.")
+
+        root = root or osp.join(self.root, 'on_disk')
+
+        in_memory_dataset = self
+        ref_data = in_memory_dataset.get(0)
+        if not isinstance(ref_data, Data):
+            raise NotImplementedError(
+                f"`{self.__class__.__name__}.to_on_disk_dataset()` is "
+                f"currently only supported on homogeneous graphs")
+
+        # Parse the schema ====================================================
+
+        schema: Dict[str, Any] = {}
+        for key, value in ref_data.to_dict().items():
+            if isinstance(value, (int, float, str)):
+                schema[key] = value.__class__
+            elif isinstance(value, Tensor) and value.dim() == 0:
+                schema[key] = dict(dtype=value.dtype, size=(-1, ))
+            elif isinstance(value, Tensor):
+                size = list(value.size())
+                size[ref_data.__cat_dim__(key, value)] = -1
+                schema[key] = dict(dtype=value.dtype, size=tuple(size))
+            else:
+                schema[key] = object
+
+        # Create the on-disk dataset ==========================================
+
+        class OnDiskDataset(torch_geometric.data.OnDiskDataset):
+            def __init__(
+                self,
+                root: str,
+                transform: Optional[Callable] = None,
+            ):
+                super().__init__(
+                    root=root,
+                    transform=transform,
+                    backend=backend,
+                    schema=schema,
+                )
+
+            def process(self):
+                _iter = [
+                    in_memory_dataset.get(i)
+                    for i in in_memory_dataset.indices()
+                ]
+                if log:  # pragma: no cover
+                    _iter = tqdm(_iter, desc='Converting to OnDiskDataset')
+
+                data_list: List[Data] = []
+                for i, data in enumerate(_iter):
+                    data_list.append(data)
+                    if i + 1 == len(in_memory_dataset) or (i + 1) % 1000 == 0:
+                        self.extend(data_list)
+                        data_list = []
+
+            def serialize(self, data: Data) -> Dict[str, Any]:
+                return data.to_dict()
+
+            def deserialize(self, data: Dict[str, Any]) -> Data:
+                return Data.from_dict(data)
+
+            def __repr__(self) -> str:
+                arg_repr = str(len(self)) if len(self) > 1 else ''
+                return (f'OnDisk{in_memory_dataset.__class__.__name__}('
+                        f'{arg_repr})')
+
+        return OnDiskDataset(root, transform=in_memory_dataset.transform)
+
     @property
     def data(self) -> Any:
         msg1 = ("It is not recommended to directly access the internal "
diff --git a/torch_geometric/data/on_disk_dataset.py b/torch_geometric/data/on_disk_dataset.py
index 3641aec59..90b4e772f 100644
--- a/torch_geometric/data/on_disk_dataset.py
+++ b/torch_geometric/data/on_disk_dataset.py
@@ -1,5 +1,4 @@
 import os
-from abc import ABC
 from typing import Any, Callable, Iterable, List, Optional, Union
 
 from torch import Tensor
@@ -10,7 +9,7 @@ from torch_geometric.data.database import Schema
 from torch_geometric.data.dataset import Dataset
 
 
-class OnDiskDataset(Dataset, ABC):
+class OnDiskDataset(Dataset):
     r"""Dataset base class for creating large graph datasets which do not
     easily fit into CPU memory at once by leveraging a :class:`Database`
     backend for on-disk storage and access of data objects.
diff --git a/torch_geometric/datasets/ba_shapes.py b/torch_geometric/datasets/ba_shapes.py
index 867114d12..5be666d86 100644
--- a/torch_geometric/datasets/ba_shapes.py
+++ b/torch_geometric/datasets/ba_shapes.py
@@ -41,7 +41,7 @@ class BAShapes(InMemoryDataset):
     """
     def __init__(self, connection_distribution: str = "random",
                  transform: Optional[Callable] = None):
-        super().__init__('.', transform)
+        super().__init__(None, transform)
         assert connection_distribution in ['random', 'uniform']
 
         # Build the Barabasi-Albert graph:
diff --git a/torch_geometric/datasets/fake.py b/torch_geometric/datasets/fake.py
index b130aa628..8f1e2b6b9 100644
--- a/torch_geometric/datasets/fake.py
+++ b/torch_geometric/datasets/fake.py
@@ -53,7 +53,7 @@ class FakeDataset(InMemoryDataset):
         pre_transform: Optional[Callable] = None,
         **kwargs,
     ):
-        super().__init__('.', transform)
+        super().__init__(None, transform)
 
         if task == 'auto':
             task = 'graph' if num_graphs > 1 else 'node'
@@ -150,7 +150,7 @@ class FakeHeteroDataset(InMemoryDataset):
         pre_transform: Optional[Callable] = None,
         **kwargs,
     ):
-        super().__init__('.', transform)
+        super().__init__(None, transform)
 
         if task == 'auto':
             task = 'graph' if num_graphs > 1 else 'node'
diff --git a/torch_geometric/datasets/karate.py b/torch_geometric/datasets/karate.py
index 9aa0ec848..7da24dd3e 100644
--- a/torch_geometric/datasets/karate.py
+++ b/torch_geometric/datasets/karate.py
@@ -38,7 +38,7 @@ class KarateClub(InMemoryDataset):
           - 4
     """
     def __init__(self, transform: Optional[Callable] = None):
-        super().__init__('.', transform)
+        super().__init__(None, transform)
 
         row = [
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1,


Test patch:
diff --git a/test/data/test_dataset.py b/test/data/test_dataset.py
index ad0538557..c6ad968dd 100644
--- a/test/data/test_dataset.py
+++ b/test/data/test_dataset.py
@@ -4,7 +4,9 @@ import pytest
 import torch
 
 from torch_geometric.data import Data, HeteroData, InMemoryDataset
+from torch_geometric.datasets import KarateClub
 from torch_geometric.testing import withPackage
+from torch_geometric.transforms import BaseTransform
 from torch_geometric.typing import SparseTensor
 
 
@@ -391,3 +393,38 @@ def test_file_names_as_property_and_method():
             pass
 
     MyTestDataset()
+
+
+@withPackage('sqlite3')
+def test_to_on_disk_dataset(tmp_path):
+    class MyTransform(BaseTransform):
+        def forward(self, data: Data) -> Data:
+            data.z = 'test_str'
+            return data
+
+    in_memory_dataset = KarateClub(transform=MyTransform())
+
+    with pytest.raises(ValueError, match="root directory of 'KarateClub'"):
+        in_memory_dataset.to_on_disk_dataset()
+
+    on_disk_dataset = in_memory_dataset.to_on_disk_dataset(tmp_path, log=False)
+    assert str(on_disk_dataset) == 'OnDiskKarateClub()'
+    assert on_disk_dataset.schema == {
+        'x': dict(dtype=torch.float32, size=(-1, 34)),
+        'edge_index': dict(dtype=torch.int64, size=(2, -1)),
+        'y': dict(dtype=torch.int64, size=(-1, )),
+        'train_mask': dict(dtype=torch.bool, size=(-1, )),
+    }
+    assert in_memory_dataset.transform == on_disk_dataset.transform
+
+    data1 = in_memory_dataset[0]
+    data2 = on_disk_dataset[0]
+
+    assert len(data1) == len(data2)
+    assert torch.allclose(data1.x, data2.x)
+    assert torch.equal(data1.edge_index, data2.edge_index)
+    assert torch.equal(data1.y, data2.y)
+    assert torch.equal(data1.train_mask, data2.train_mask)
+    assert data1.z == data2.z
+
+    on_disk_dataset.close()
diff --git a/test/loader/test_dataloader.py b/test/loader/test_dataloader.py
index a314081c6..50113f380 100644
--- a/test/loader/test_dataloader.py
+++ b/test/loader/test_dataloader.py
@@ -1,7 +1,6 @@
 import multiprocessing
 import sys
 from collections import namedtuple
-from typing import Any, Dict, List
 
 import pytest
 import torch
@@ -205,44 +204,7 @@ if __name__ == '__main__':
             pass
         print(f'Done! [{time.perf_counter() - t:.4f}s]')
 
-    class OnDiskQM9(OnDiskDataset):
-        def __init__(self, root: str):
-            schema = {
-                'x': dict(dtype=torch.float32, size=(-1, 11)),
-                'edge_index': dict(dtype=torch.int64, size=(2, -1)),
-                'edge_attr': dict(dtype=torch.float, size=(-1, 8)),
-                'y': dict(dtype=torch.float, size=(1, 19)),
-                'z': dict(dtype=torch.int64, size=(-1, )),
-                'smiles': str,
-                'name': str,
-                'idx': dict(dtype=torch.int64, size=(-1, )),
-            }
-            super().__init__(root, schema=schema)
-
-        def process(self):
-            data_list: List[Data] = []
-            for i, data in enumerate(in_memory_dataset):
-                data_list.append(data)
-                if i + 1 == len(in_memory_dataset) or (i + 1) % 1000 == 0:
-                    self.extend(data_list)
-                    data_list = []
-
-        def serialize(self, data: Data) -> Dict[str, Any]:
-            return dict(
-                x=data.x,
-                edge_index=data.edge_index,
-                edge_attr=data.edge_attr,
-                y=data.y,
-                z=data.z,
-                smiles=data.smiles,
-                name=data.name,
-                idx=data.idx,
-            )
-
-        def deserialize(self, data: Dict[str, Any]) -> Data:
-            return Data.from_dict(data)
-
-    on_disk_dataset = OnDiskQM9(root='/tmp/OnDiskQM9')
+    on_disk_dataset = in_memory_dataset.to_on_disk_dataset()
     loader = DataLoader(on_disk_dataset, **kwargs)
 
     print('On-Disk Dataset:')
@@ -252,3 +214,5 @@ if __name__ == '__main__':
         for batch in loader:
             pass
         print(f'Done! [{time.perf_counter() - t:.4f}s]')
+
+    on_disk_dataset.close()