From b24871569213f8d4da97403c0d7474767dceeaeb Mon Sep 17 00:00:00 2001 From: Jeffrey Qiu <77599736+hongjie-qiu@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:41:10 +0800 Subject: [PATCH] Fix GridPatchDataset cache edge cases (#9100) Signed-off-by: Jeffrey Qiu <77599736+hongjie-qiu@users.noreply.github.com> --- monai/data/grid_dataset.py | 17 ++++++++------- tests/data/test_grid_dataset.py | 37 +++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/monai/data/grid_dataset.py b/monai/data/grid_dataset.py index a4860cfcad5..00dbb370b50 100644 --- a/monai/data/grid_dataset.py +++ b/monai/data/grid_dataset.py @@ -257,6 +257,8 @@ def __init__( self.first_random = self.patch_transform.get_index_of_first( lambda t: isinstance(t, RandomizableTrait) or not isinstance(t, Transform) ) + if self.first_random is None: + self.first_random = len(self.patch_transform.transforms) if self.cache: if isinstance(data, Iterator): @@ -279,7 +281,11 @@ def set_data(self, data: Sequence) -> None: self.cache_num = min(int(self.set_num), int(len(mapping) * self.set_rate), len(mapping)) self._hash_keys = list(mapping)[: self.cache_num] indices = list(mapping.values())[: self.cache_num] - self._cache, self._cache_other = zip(*self._fill_cache(indices)) # type: ignore + cache_items = self._fill_cache(indices) + if cache_items: + self._cache, self._cache_other = zip(*cache_items) # type: ignore + else: + self._cache, self._cache_other = [], [] def _fill_cache(self, indices=None) -> list: """ @@ -339,12 +345,9 @@ def _generate_patches(self, src, **apply_args): def __iter__(self): if self.cache: - cache_index = None for image in super().__iter__(): key = self.hash_func(image) - if key in self._hash_keys: - # if existing in cache, try to get the index in cache - cache_index = self._hash_keys.index(key) + cache_index = self._hash_keys.index(key) if key in self._hash_keys else None if cache_index is None: # no cache for this index, execute all the transforms directly yield from self._generate_patches(self.patch_iter(image)) @@ -354,11 +357,11 @@ def __iter__(self): "Cache buffer is not initialized, please call `set_data()` before epoch begins." ) data = self._cache[cache_index] - other = self._cache_other[cache_index] # load data from cache and execute from the first random transform data = deepcopy(data) if self.copy_cache else data - yield from self._generate_patches(zip(data, other), start=self.first_random) + cached_patches = zip(data, self._cache_other[cache_index]) if self.with_coordinates else zip(data) + yield from self._generate_patches(cached_patches, start=self.first_random) else: for image in super().__iter__(): yield from self._generate_patches(self.patch_iter(image)) diff --git a/tests/data/test_grid_dataset.py b/tests/data/test_grid_dataset.py index 0b44564834c..8053027a14c 100644 --- a/tests/data/test_grid_dataset.py +++ b/tests/data/test_grid_dataset.py @@ -200,6 +200,7 @@ def test_set_data(self): np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2)) np.testing.assert_allclose(item[0], np.array([[[[81, 91], [121, 131]]], [[[101, 111], [141, 151]]]]), rtol=1e-4) np.testing.assert_allclose(item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5) + # simulate another epoch, the cache content should not be modified for item in DataLoader(dataset, batch_size=2, shuffle=False, num_workers=num_workers): np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2)) @@ -217,6 +218,42 @@ def test_set_data(self): ) np.testing.assert_allclose(item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5) + def test_partial_cache_preserves_uncached_items(self): + dataset = GridPatchDataset( + data=[[1], [2]], patch_iter=identity_generator, cache=True, cache_rate=0.5, progress=False + ) + + self.assertEqual(list(dataset), [(1, 0), (2, 0)]) + + def test_cache_without_coordinates(self): + dataset = GridPatchDataset( + data=[[1, 2]], patch_iter=identity_generator, with_coordinates=False, cache=True, progress=False + ) + + self.assertEqual(list(dataset), [1, 2]) + + def test_cache_with_deterministic_transform(self): + from monai.transforms import Lambda + + dataset = GridPatchDataset( + data=[[1, 2]], + patch_iter=identity_generator, + transform=Lambda(lambda x: x + 100), + cache=True, + progress=False, + ) + + self.assertEqual(list(dataset), [(101, 0), (102, 1)]) + + def test_zero_sized_cache(self): + for cache_kwargs in ({"cache_rate": 0.0}, {"cache_num": 0}): + with self.subTest(**cache_kwargs): + dataset = GridPatchDataset( + data=[[1], [2]], patch_iter=identity_generator, cache=True, progress=False, **cache_kwargs + ) + + self.assertEqual(list(dataset), [(1, 0), (2, 0)]) + if __name__ == "__main__": unittest.main()