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,11 @@ 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.
+                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 +299,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 +337,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 +446,22 @@ 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] = self._cachecheck(self.data[index])
+        # 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,71 @@ 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_metatensor_loading(self):
         """
         Thorough test of metadata loading correctly with MetaTensor. This will store a MetaTensor with safe object types