From 749a5182fce4ef46b973f90d4f93924de7edffa3 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Mon, 12 Jan 2026 10:55:07 +0000 Subject: [PATCH 1/6] Add in_memory option to PersistentDataset for hybrid caching - Add `in_memory` parameter to PersistentDataset that combines persistent disk storage with RAM caching for faster subsequent access - Add `memory_cache_size` property for inspecting cache state - Compute cache key once in _cachecheck() to avoid redundant hash computation - Clear memory cache when set_data() is called - Works with or without cache_dir (pure RAM cache mode) Addresses: https://github.com/Project-MONAI/MONAI/issues/6753 Signed-off-by: Soumya Snigdha Kundu --- monai/data/dataset.py | 35 ++++++++- tests/data/test_persistentdataset.py | 106 +++++++++++++++++++++++++++ 2 files changed, 137 insertions(+), 4 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 066cec41b7f..1db4ce9558a 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -232,6 +232,7 @@ def __init__( reset_ops_id: bool = True, track_meta: bool = False, weights_only: bool = True, + in_memory: bool = False, ) -> None: """ Args: @@ -273,6 +274,10 @@ def __init__( other safe objects. Setting this to `False` is required for loading `MetaTensor` objects saved with `track_meta=True`, however this creates the possibility of remote code execution through `torch.load` so be aware of the security implications of doing so. + in_memory: if `True`, keep the pre-processed data in an in-memory dictionary after first access. + This combines the benefits of persistent storage (data survives restarts) with faster RAM access. + When data is accessed, it is first loaded from disk cache and then stored in memory. + Default to `False`. Raises: ValueError: When both `track_meta=True` and `weights_only=True`, since this combination @@ -299,6 +304,13 @@ def __init__( ) self.track_meta = track_meta self.weights_only = weights_only + self.in_memory = in_memory + self._memory_cache: dict[str, Any] = {} + + @property + def memory_cache_size(self) -> int: + """Return 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 @@ -326,6 +338,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) @@ -389,14 +402,24 @@ def _cachecheck(self, item_transformed): """ hashfile = None + # compute cache key once for both disk and memory caching + data_item_md5 = self.hash_func(item_transformed).decode("utf-8") + data_item_md5 += self.transform_hash + cache_key = f"{data_item_md5}.pt" + if self.cache_dir is not None: - data_item_md5 = self.hash_func(item_transformed).decode("utf-8") - data_item_md5 += self.transform_hash - hashfile = self.cache_dir / f"{data_item_md5}.pt" + hashfile = self.cache_dir / cache_key + + # check in-memory cache first + if self.in_memory and cache_key in self._memory_cache: + return self._memory_cache[cache_key] if hashfile is not None and hashfile.is_file(): # cache hit try: - return torch.load(hashfile, weights_only=self.weights_only) + _item_transformed = torch.load(hashfile, weights_only=self.weights_only) + if self.in_memory: + self._memory_cache[cache_key] = _item_transformed + return _item_transformed except PermissionError as e: if sys.platform != "win32": raise e @@ -409,6 +432,8 @@ def _cachecheck(self, item_transformed): _item_transformed = self._pre_transform(deepcopy(item_transformed)) # keep the original hashed if hashfile is None: + if self.in_memory: + self._memory_cache[cache_key] = _item_transformed return _item_transformed try: # NOTE: Writing to a temporary directory and then using a nearly atomic rename operation @@ -431,6 +456,8 @@ def _cachecheck(self, item_transformed): pass except PermissionError: # project-monai/monai issue #3613 pass + if self.in_memory: + self._memory_cache[cache_key] = _item_transformed return _item_transformed def _transform(self, index: int): diff --git a/tests/data/test_persistentdataset.py b/tests/data/test_persistentdataset.py index ca62cdb1840..99129045024 100644 --- a/tests/data/test_persistentdataset.py +++ b/tests/data/test_persistentdataset.py @@ -15,6 +15,7 @@ import os import tempfile import unittest +from pathlib import Path import nibabel as nib import numpy as np @@ -200,6 +201,111 @@ 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): + """Test in_memory caching feature that combines persistent storage with RAM caching.""" + items = [[list(range(i))] for i in range(5)] + + with tempfile.TemporaryDirectory() as tempdir: + # First, create the persistent cache + ds1 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=False) + # Access all items to populate disk cache + _ = list(ds1) + + # Now create a new dataset with in_memory=True + ds2 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True) + + # Memory cache should be empty initially + self.assertEqual(ds2.memory_cache_size, 0) + + # Access items - they should be loaded from disk and cached in memory + _ = ds2[0] + self.assertEqual(ds2.memory_cache_size, 1) + + _ = ds2[1] + self.assertEqual(ds2.memory_cache_size, 2) + + # Access all items + _ = list(ds2) + self.assertEqual(ds2.memory_cache_size, 5) + + # Accessing same item again should use memory cache (same result) + result1 = ds2[0] + result2 = ds2[0] + self.assertEqual(result1, result2) + + # Test set_data clears in-memory cache + ds2.set_data(items[:3]) + self.assertEqual(ds2.memory_cache_size, 0) + + 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_automatic_hybrid_caching(self): + """ + Test that in_memory=True provides automatic hybrid caching: + - ALL samples automatically persist to disk + - ALL samples automatically cache to RAM after first access + - No manual specification of which samples go where (unlike torchdatasets) + - Simulates restart scenario: disk cache survives, RAM cache rebuilds automatically + """ + items = [[list(range(i))] for i in range(5)] + + with tempfile.TemporaryDirectory() as tempdir: + # === First "session": populate both disk and RAM cache === + ds1 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True) + + # Access all items - should automatically cache to BOTH disk AND RAM + for i in range(len(items)): + _ = ds1[i] + + # Verify: ALL samples are in RAM (automatic, no manual specification) + self.assertEqual(ds1.memory_cache_size, 5) + + # Verify: ALL samples are on disk (count .pt files) + cache_files = list(Path(tempdir).glob("*.pt")) + self.assertEqual(len(cache_files), 5) + + # === Simulate "restart": new dataset instance, same cache_dir === + # This is the key benefit over CacheDataset - disk cache survives restart + ds2 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True) + + # RAM cache starts empty (simulating fresh process) + self.assertEqual(ds2.memory_cache_size, 0) + + # Access all items - should load from disk and automatically cache to RAM + results = [ds2[i] for i in range(len(items))] + + # Verify: ALL samples now in RAM again (automatic rebuild from disk) + self.assertEqual(ds2.memory_cache_size, 5) + + # Verify: Results are correct + for i, result in enumerate(results): + self.assertEqual(result, [list(range(i))]) + + # === Verify RAM cache provides fast repeated access === + # Accessing same items again should hit RAM cache (same objects) + for i in range(len(items)): + result1 = ds2[i] + result2 = ds2[i] + # Should return equivalent results + self.assertEqual(result1, result2) + + # RAM cache size unchanged (no duplicate entries) + self.assertEqual(ds2.memory_cache_size, 5) + if __name__ == "__main__": unittest.main() From 1d873cf162875808f1ac24342bd1dc56bd4272dd Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Wed, 28 Jan 2026 03:21:30 +0000 Subject: [PATCH 2/6] Fix type consistency in memory cache and correct test assertions - Store tensor-converted data in memory cache to match disk cache types - Fix test_automatic_hybrid_caching assertions to account for _InplaceXform (from coderabbit) Signed-off-by: Soumya Snigdha Kundu --- monai/data/dataset.py | 6 ++++-- tests/data/test_persistentdataset.py | 8 ++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 1db4ce9558a..c6807265a2f 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -435,6 +435,8 @@ def _cachecheck(self, item_transformed): if self.in_memory: self._memory_cache[cache_key] = _item_transformed return _item_transformed + # Convert to tensor for disk storage (and memory cache consistency) + _item_converted = convert_to_tensor(_item_transformed, convert_numeric=False, track_meta=self.track_meta) try: # NOTE: Writing to a temporary directory and then using a nearly atomic rename operation # to make the cache more robust to manual killing of parent process @@ -442,7 +444,7 @@ def _cachecheck(self, item_transformed): with tempfile.TemporaryDirectory() as tmpdirname: temp_hash_file = Path(tmpdirname) / hashfile.name torch.save( - obj=convert_to_tensor(_item_transformed, convert_numeric=False, track_meta=self.track_meta), + obj=_item_converted, f=temp_hash_file, pickle_module=look_up_option(self.pickle_module, SUPPORTED_PICKLE_MOD), pickle_protocol=self.pickle_protocol, @@ -457,7 +459,7 @@ def _cachecheck(self, item_transformed): except PermissionError: # project-monai/monai issue #3613 pass if self.in_memory: - self._memory_cache[cache_key] = _item_transformed + self._memory_cache[cache_key] = _item_converted return _item_transformed def _transform(self, index: int): diff --git a/tests/data/test_persistentdataset.py b/tests/data/test_persistentdataset.py index 99129045024..2aa67be555a 100644 --- a/tests/data/test_persistentdataset.py +++ b/tests/data/test_persistentdataset.py @@ -291,9 +291,13 @@ def test_automatic_hybrid_caching(self): # Verify: ALL samples now in RAM again (automatic rebuild from disk) self.assertEqual(ds2.memory_cache_size, 5) - # Verify: Results are correct + # Verify: Results are correct (transformed by _InplaceXform) for i, result in enumerate(results): - self.assertEqual(result, [list(range(i))]) + if i == 0: + expected = [[1]] # empty list -> append 1 + else: + expected = [[np.pi] + list(range(1, i))] # data[0] = 0 + np.pi + self.assertEqual(result, expected) # === Verify RAM cache provides fast repeated access === # Accessing same items again should hit RAM cache (same objects) From 152f7a098f7395f2febfae04b9c431068178a9bf Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 30 Jul 2026 15:00:47 +0100 Subject: [PATCH 3/6] Simplify in_memory cache: memoize in _transform by index Signed-off-by: Soumya Snigdha Kundu --- monai/data/dataset.py | 42 +++++------- tests/data/test_persistentdataset.py | 98 ++++++---------------------- 2 files changed, 35 insertions(+), 105 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 1f79bf11950..5b962a633d1 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -277,9 +277,10 @@ 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`, keep the pre-processed data in an in-memory dictionary after first access. - This combines the benefits of persistent storage (data survives restarts) with faster RAM access. - When data is accessed, it is first loaded from disk cache and then stored in memory. + 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) @@ -299,7 +300,7 @@ def __init__( self.track_meta = track_meta self.weights_only = weights_only self.in_memory = in_memory - self._memory_cache: dict[str, Any] = {} + self._memory_cache: dict[int, Any] = {} @property def memory_cache_size(self) -> int: @@ -396,24 +397,14 @@ def _cachecheck(self, item_transformed): """ hashfile = None - # compute cache key once for both disk and memory caching - data_item_md5 = self.hash_func(item_transformed).decode("utf-8") - data_item_md5 += self.transform_hash - cache_key = f"{data_item_md5}.pt" - if self.cache_dir is not None: - hashfile = self.cache_dir / cache_key - - # check in-memory cache first - if self.in_memory and cache_key in self._memory_cache: - return self._memory_cache[cache_key] + data_item_hash = self.hash_func(item_transformed).decode("utf-8") + data_item_hash += self.transform_hash + hashfile = self.cache_dir / f"{data_item_hash}.pt" if hashfile is not None and hashfile.is_file(): # cache hit try: - _item_transformed = torch.load(hashfile, weights_only=self.weights_only) - if self.in_memory: - self._memory_cache[cache_key] = _item_transformed - return _item_transformed + return torch.load(hashfile, weights_only=self.weights_only) except PermissionError as e: if sys.platform != "win32": raise e @@ -426,11 +417,7 @@ def _cachecheck(self, item_transformed): _item_transformed = self._pre_transform(deepcopy(item_transformed)) # keep the original hashed if hashfile is None: - if self.in_memory: - self._memory_cache[cache_key] = _item_transformed return _item_transformed - # Convert to tensor for disk storage (and memory cache consistency) - _item_converted = convert_to_tensor(_item_transformed, convert_numeric=False, track_meta=self.track_meta) try: # NOTE: Writing to a temporary directory and then using a nearly atomic rename operation # to make the cache more robust to manual killing of parent process @@ -438,7 +425,7 @@ def _cachecheck(self, item_transformed): with tempfile.TemporaryDirectory() as tmpdirname: temp_hash_file = Path(tmpdirname) / hashfile.name torch.save( - obj=_item_converted, + obj=convert_to_tensor(_item_transformed, convert_numeric=False, track_meta=self.track_meta), f=temp_hash_file, pickle_module=look_up_option(self.pickle_module, SUPPORTED_PICKLE_MOD), pickle_protocol=self.pickle_protocol, @@ -452,13 +439,14 @@ def _cachecheck(self, item_transformed): pass except PermissionError: # project-monai/monai issue #3613 pass - if self.in_memory: - self._memory_cache[cache_key] = _item_converted return _item_transformed def _transform(self, index: int): - pre_random_item = self._cachecheck(self.data[index]) - return self._post_transform(pre_random_item) + 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]) + return self._post_transform(self._memory_cache[index]) class CacheNTransDataset(PersistentDataset): diff --git a/tests/data/test_persistentdataset.py b/tests/data/test_persistentdataset.py index 2a2f51e88fd..3f4c00a0abc 100644 --- a/tests/data/test_persistentdataset.py +++ b/tests/data/test_persistentdataset.py @@ -204,38 +204,37 @@ def test_track_meta_and_weights_only(self, track_meta, weights_only, expected_er self.assertIsInstance(im, expected_type) def test_in_memory_cache(self): - """Test in_memory caching feature that combines persistent storage with RAM caching.""" + """`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, create the persistent cache - ds1 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=False) - # Access all items to populate disk cache + # 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) - # Now create a new dataset with in_memory=True + # 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) - - # Memory cache should be empty initially self.assertEqual(ds2.memory_cache_size, 0) - # Access items - they should be loaded from disk and cached in memory - _ = ds2[0] - self.assertEqual(ds2.memory_cache_size, 1) - - _ = ds2[1] - self.assertEqual(ds2.memory_cache_size, 2) - - # Access all items - _ = list(ds2) + 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) - # Accessing same item again should use memory cache (same result) - result1 = ds2[0] - result2 = ds2[0] - self.assertEqual(result1, result2) + # repeated access is served from RAM without adding entries + self.assertEqual(ds2[0], results[0]) + self.assertEqual(ds2.memory_cache_size, 5) - # Test set_data clears in-memory cache + # set_data clears the in-memory cache ds2.set_data(items[:3]) self.assertEqual(ds2.memory_cache_size, 0) @@ -255,63 +254,6 @@ def test_in_memory_without_cache_dir(self): _ = list(ds) self.assertEqual(ds.memory_cache_size, 3) - def test_automatic_hybrid_caching(self): - """ - Test that in_memory=True provides automatic hybrid caching: - - ALL samples automatically persist to disk - - ALL samples automatically cache to RAM after first access - - No manual specification of which samples go where (unlike torchdatasets) - - Simulates restart scenario: disk cache survives, RAM cache rebuilds automatically - """ - items = [[list(range(i))] for i in range(5)] - - with tempfile.TemporaryDirectory() as tempdir: - # === First "session": populate both disk and RAM cache === - ds1 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True) - - # Access all items - should automatically cache to BOTH disk AND RAM - for i in range(len(items)): - _ = ds1[i] - - # Verify: ALL samples are in RAM (automatic, no manual specification) - self.assertEqual(ds1.memory_cache_size, 5) - - # Verify: ALL samples are on disk (count .pt files) - cache_files = list(Path(tempdir).glob("*.pt")) - self.assertEqual(len(cache_files), 5) - - # === Simulate "restart": new dataset instance, same cache_dir === - # This is the key benefit over CacheDataset - disk cache survives restart - ds2 = PersistentDataset(data=items, transform=_InplaceXform(), cache_dir=tempdir, in_memory=True) - - # RAM cache starts empty (simulating fresh process) - self.assertEqual(ds2.memory_cache_size, 0) - - # Access all items - should load from disk and automatically cache to RAM - results = [ds2[i] for i in range(len(items))] - - # Verify: ALL samples now in RAM again (automatic rebuild from disk) - self.assertEqual(ds2.memory_cache_size, 5) - - # Verify: Results are correct (transformed by _InplaceXform) - for i, result in enumerate(results): - if i == 0: - expected = [[1]] # empty list -> append 1 - else: - expected = [[np.pi] + list(range(1, i))] # data[0] = 0 + np.pi - self.assertEqual(result, expected) - - # === Verify RAM cache provides fast repeated access === - # Accessing same items again should hit RAM cache (same objects) - for i in range(len(items)): - result1 = ds2[i] - result2 = ds2[i] - # Should return equivalent results - self.assertEqual(result1, result2) - - # RAM cache size unchanged (no duplicate entries) - self.assertEqual(ds2.memory_cache_size, 5) - def test_metatensor_loading(self): """ Thorough test of metadata loading correctly with MetaTensor. This will store a MetaTensor with safe object types From efcfa46eb5e768220cf6f1b8510f380f49a4d25a Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 30 Jul 2026 15:11:49 +0100 Subject: [PATCH 4/6] Copy in-memory cached items before post-transform Signed-off-by: Soumya Snigdha Kundu --- monai/data/dataset.py | 19 +++++++++++++++++-- tests/data/test_persistentdataset.py | 14 ++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 5b962a633d1..2ab5d9411fe 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -304,7 +304,11 @@ def __init__( @property def memory_cache_size(self) -> int: - """Return the number of items currently stored in the in-memory cache.""" + """ + 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]): @@ -442,11 +446,22 @@ def _cachecheck(self, item_transformed): return _item_transformed def _transform(self, index: int): + """ + 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]) - return self._post_transform(self._memory_cache[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): diff --git a/tests/data/test_persistentdataset.py b/tests/data/test_persistentdataset.py index 3f4c00a0abc..90a37e77192 100644 --- a/tests/data/test_persistentdataset.py +++ b/tests/data/test_persistentdataset.py @@ -238,6 +238,20 @@ def test_in_memory_cache(self): 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)] From 876b79988691cb2e1b76cf18f058588144244c19 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Tue, 4 Aug 2026 12:30:48 +0100 Subject: [PATCH 5/6] Normalize in-memory PersistentDataset cache to match disk round-trip When in_memory=True, the first-computed pre-random value was cached verbatim and reused for all later reads, so it never went through the convert_to_tensor(..., track_meta=self.track_meta) normalization that the disk-cache round-trip applies on every read. This made the returned type (e.g. numpy.ndarray vs torch.Tensor/MetaTensor) silently depend on in_memory, including when track_meta=True. Cache the value returned by the same convert_to_tensor call used before torch.save so the RAM-cached entry matches a disk-cache hit, and add a regression test. Flagged by @garciadias during review. Signed-off-by: Soumya Snigdha Kundu --- monai/data/dataset.py | 4 +++- tests/data/test_persistentdataset.py | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 2ab5d9411fe..a6e31c9beed 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -459,7 +459,9 @@ def _transform(self, index: int): 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]) + 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])) diff --git a/tests/data/test_persistentdataset.py b/tests/data/test_persistentdataset.py index 90a37e77192..c1eb67fcb28 100644 --- a/tests/data/test_persistentdataset.py +++ b/tests/data/test_persistentdataset.py @@ -268,6 +268,30 @@ def test_in_memory_without_cache_dir(self): _ = 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 From 247a1d877f8e068fe6e76154e74f196e8f0989bf Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Tue, 4 Aug 2026 12:39:26 +0100 Subject: [PATCH 6/6] Document persistent_workers caveat for in_memory PersistentDataset Note in the in_memory docstring that the RAM cache is per-process, so DataLoader(num_workers>0, persistent_workers=False) rebuilds it every epoch and loses most of the benefit; persistent_workers=True or num_workers=0 is needed to retain it across epochs. Flagged by @garciadias during review. Signed-off-by: Soumya Snigdha Kundu --- monai/data/dataset.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index a6e31c9beed..d7045d15cfc 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -281,6 +281,9 @@ def __init__( 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)