Fix GridPatchDataset cache edge cases (#9100)

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))

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()