Add in_memory option to PersistentDataset for hybrid caching

monai/data/dataset.py

@@ -236,6 +236,7 @@ def __init__(
         reset_ops_id: bool = True,
         track_meta: bool = False,
         weights_only: bool = True,
+        in_memory: bool = False,
     ) -> None:
         """
         Args:
@@ -276,6 +277,14 @@ def __init__(
                 Setting to `False` should only be done if it's absolutely necessary to load unsafe pickled data,
                 eg. MetaTensor objects with unsafe objects in their metadata. Users must verify the safety of the data
                 they intend to load before doing so.
+            in_memory: if `True`, also keep the pre-processed data in RAM after first access, so that later
+                epochs skip the disk cache entirely. This combines the benefits of persistent storage (data
+                survives restarts) with faster RAM access. Note the cache is unbounded, so the whole dataset
+                is eventually held in memory; use `CacheDataset` or `SmartCacheDataset` if that does not fit.
+                The RAM cache is per-process, so with `DataLoader(num_workers>0, persistent_workers=False)`
+                the cache is rebuilt every epoch; use `persistent_workers=True` (or `num_workers=0`) to retain
+                it across epochs, otherwise most of the benefit is lost.
+                Default to `False`.
         """
         super().__init__(data=data, transform=transform)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else None
@@ -293,6 +302,17 @@ def __init__(
         self.reset_ops_id = reset_ops_id
         self.track_meta = track_meta
         self.weights_only = weights_only
+        self.in_memory = in_memory
+        self._memory_cache: dict[int, Any] = {}
+
+    @property
+    def memory_cache_size(self) -> int:
+        """
+        Returns:
+            The number of items currently stored in the in-memory cache.
+
+        """
+        return len(self._memory_cache)
 
     def set_transform_hash(self, hash_xform_func: Callable[..., bytes]):
         """Get hashable transforms, and then hash them. Hashable transforms
@@ -320,6 +340,7 @@ def set_data(self, data: Sequence):
 
         """
         self.data = data
+        self._memory_cache = {}
         if self.cache_dir is not None and self.cache_dir.exists():
             shutil.rmtree(self.cache_dir, ignore_errors=True)
             self.cache_dir.mkdir(parents=True, exist_ok=True)
@@ -428,8 +449,24 @@ def _cachecheck(self, item_transformed):
         return _item_transformed
 
     def _transform(self, index: int):
-        pre_random_item = self._cachecheck(self.data[index])
-        return self._post_transform(pre_random_item)
+        """
+        Fetch the pre-random-transform item for `index` and apply the random transforms to it.
+
+        Args:
+            index: index of the item in `self.data`.
+
+        Returns:
+            The fully transformed data element.
+
+        """
+        if not self.in_memory:
+            return self._post_transform(self._cachecheck(self.data[index]))
+        if index not in self._memory_cache:
+            self._memory_cache[index] = convert_to_tensor(
+                self._cachecheck(self.data[index]), convert_numeric=False, track_meta=self.track_meta
+            )
+        # copy so that the random transforms, or the caller, cannot mutate the cached item
+        return self._post_transform(deepcopy(self._memory_cache[index]))
 
 
 class CacheNTransDataset(PersistentDataset):

tests/data/test_persistentdataset.py

@@ -203,6 +203,95 @@ def test_track_meta_and_weights_only(self, track_meta, weights_only, expected_er
                 im = test_dataset[0]["image"]
                 self.assertIsInstance(im, expected_type)
 
+    def test_in_memory_cache(self):
+        """`in_memory=True` caches to RAM on top of the disk cache, and rebuilds that RAM cache after a restart."""
+        items = [[list(range(i))] for i in range(5)]
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            # first "session": every accessed item is written to disk and kept in RAM
+            ds1 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True)
+            self.assertEqual(ds1.memory_cache_size, 0)
+
+            _ = ds1[0]
+            self.assertEqual(ds1.memory_cache_size, 1)
+
+            _ = list(ds1)
+            self.assertEqual(ds1.memory_cache_size, 5)
+            self.assertEqual(len(list(Path(tempdir).glob("*.pt"))), 5)
+
+            # simulate a restart: the disk cache survives, the RAM cache is rebuilt from it
+            ds2 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True)
+            self.assertEqual(ds2.memory_cache_size, 0)
+
+            results = [ds2[i] for i in range(len(items))]
+            self.assertEqual(ds2.memory_cache_size, 5)
+            for i, result in enumerate(results):
+                # data[0] = 0 + np.pi, except for the empty item which gets 1 appended
+                expected = [[1]] if i == 0 else [[np.pi] + list(range(1, i))]
+                self.assertEqual(result, expected)
+
+            # repeated access is served from RAM without adding entries
+            self.assertEqual(ds2[0], results[0])
+            self.assertEqual(ds2.memory_cache_size, 5)
+
+            # set_data clears the in-memory cache
+            ds2.set_data(items[:3])
+            self.assertEqual(ds2.memory_cache_size, 0)
+
+    def test_in_memory_mutation_isolation(self):
+        """Mutating a returned item must not corrupt the RAM-cached copy used by later reads."""
+        items = [[list(range(i))] for i in range(3)]
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            for cache_dir in (None, tempdir):
+                with self.subTest(cache_dir=cache_dir):
+                    ds = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=cache_dir, in_memory=True)
+                    first = ds[2]
+                    self.assertEqual(first, [[np.pi, 1]])
+
+                    first[0].append(999)  # caller mutates the item it was handed
+                    self.assertEqual(ds[2], [[np.pi, 1]])
+
+    def test_in_memory_without_cache_dir(self):
+        """Test in_memory caching works even without a cache_dir (pure RAM cache)."""
+        items = [[list(range(i))] for i in range(3)]
+
+        ds = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=None, in_memory=True)
+
+        # Memory cache should be empty initially
+        self.assertEqual(ds.memory_cache_size, 0)
+
+        # Access items - they should be cached in memory
+        _ = ds[0]
+        self.assertEqual(ds.memory_cache_size, 1)
+
+        _ = list(ds)
+        self.assertEqual(ds.memory_cache_size, 3)
+
+    def test_in_memory_type_consistency(self):
+        """The RAM cache applies the same convert_to_tensor(..., track_meta=...) normalization as the disk
+        round-trip, so repeated reads of the same index return a tensor type regardless of `in_memory`."""
+
+        class _ToNumpyXform(Transform):
+            def __call__(self, data):
+                data["image"] = np.zeros((2, 2), dtype=np.float32)
+                return data
+
+        data = [dict() for _ in range(2)]
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            cache_dirs = {"memory": None, "disk": tempdir}
+            for label, cache_dir in cache_dirs.items():
+                with self.subTest(cache_dir=label):
+                    ds = PersistentDataset(
+                        data=data, transform=_ToNumpyXform(), cache_dir=cache_dir, in_memory=True, track_meta=True
+                    )
+                    # the pre-random transform only ever produces numpy arrays; with in_memory enabled the
+                    # normalized type must match the disk-cache round-trip on every read (not just the first)
+                    types = [type(ds[0]["image"]) for _ in range(3)]
+                    self.assertTrue(all(t is types[0] for t in types), types)
+                    self.assertIsInstance(ds[0]["image"], torch.Tensor, types)
+
     def test_metatensor_loading(self):
         """
         Thorough test of metadata loading correctly with MetaTensor. This will store a MetaTensor with safe object types