Permanent dataset hashes transforms (#4633)

* permanent dataset hashes transforms

Signed-off-by: Richard Brown <[email protected]>

* fixes

Signed-off-by: Richard Brown <[email protected]>

* documentation

Signed-off-by: Richard Brown <[email protected]>

* default to no hash

Signed-off-by: Wenqi Li <[email protected]>

Co-authored-by: Wenqi Li <[email protected]>

Author: rijobro

Diff for: monai/data/dataset.py

@@ -191,9 +191,10 @@ class PersistentDataset(Dataset):
     Note:
         The input data must be a list of file paths and will hash them as cache keys.
 
-        When loading persistent cache content, it can't guarantee the cached data matches current
-        transform chain, so please make sure to use exactly the same non-random transforms and the
-        args as the cache content, otherwise, it may cause unexpected errors.
+        The filenames of the cached files also try to contain the hash of the transforms. In this
+        fashion, `PersistentDataset` should be robust to changes in transforms. This, however, is
+        not guaranteed, so caution should be used when modifying transforms to avoid unexpected
+        errors. If in doubt, it is advisable to clear the cache directory.
 
     """
 
@@ -205,6 +206,7 @@ def __init__(
         hash_func: Callable[..., bytes] = pickle_hashing,
         pickle_module: str = "pickle",
         pickle_protocol: int = DEFAULT_PROTOCOL,
+        hash_transform: Optional[Callable[..., bytes]] = None,
     ) -> None:
         """
         Args:
@@ -232,6 +234,9 @@ def __init__(
             pickle_protocol: can be specified to override the default protocol, default to `2`.
                 this arg is used by `torch.save`, for more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
+            hash_transform: a callable to compute hash from the transform information when caching.
+                This may reduce errors due to transforms changing during experiments. Default to None (no hash).
+                Other options are `pickle_hashing` and `json_hashing` functions from `monai.data.utils`.
 
         """
         if not isinstance(transform, Compose):
@@ -246,6 +251,29 @@ def __init__(
                 self.cache_dir.mkdir(parents=True, exist_ok=True)
             if not self.cache_dir.is_dir():
                 raise ValueError("cache_dir must be a directory.")
+        self.transform_hash = ""
+        if hash_transform is not None:
+            self.set_transform_hash(hash_transform)
+
+    def set_transform_hash(self, hash_xform_func):
+        """Get hashable transforms, and then hash them. Hashable transforms
+        are deterministic transforms that inherit from `Transform`. We stop
+        at the first non-deterministic transform, or first that does not
+        inherit from MONAI's `Transform` class."""
+        hashable_transforms = []
+        for _tr in self.transform.flatten().transforms:
+            if isinstance(_tr, Randomizable) or not isinstance(_tr, Transform):
+                break
+            hashable_transforms.append(_tr)
+        # Try to hash. Fall back to a hash of their names
+        try:
+            self.transform_hash = hash_xform_func(hashable_transforms)
+        except TypeError as te:
+            if "is not JSON serializable" not in str(te):
+                raise te
+            names = "".join(tr.__class__.__name__ for tr in hashable_transforms)
+            self.transform_hash = hash_xform_func(names)
+        self.transform_hash = self.transform_hash.decode("utf-8")
 
     def set_data(self, data: Sequence):
         """
@@ -325,6 +353,7 @@ def _cachecheck(self, item_transformed):
         hashfile = None
         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"
 
         if hashfile is not None and hashfile.is_file():  # cache hit

Diff for: tests/test_persistentdataset.py

@@ -19,7 +19,7 @@
 from parameterized import parameterized
 
 from monai.data import PersistentDataset, json_hashing
-from monai.transforms import Compose, LoadImaged, SimulateDelayd, Transform
+from monai.transforms import Compose, Flip, Identity, LoadImaged, SimulateDelayd, Transform
 
 TEST_CASE_1 = [
     Compose(
@@ -77,7 +77,7 @@ def test_cache(self):
 
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
     def test_shape(self, transform, expected_shape):
-        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
+        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]).astype(float), np.eye(4))
         with tempfile.TemporaryDirectory() as tempdir:
             nib.save(test_image, os.path.join(tempdir, "test_image1.nii.gz"))
             nib.save(test_image, os.path.join(tempdir, "test_label1.nii.gz"))
@@ -150,6 +150,19 @@ def test_shape(self, transform, expected_shape):
                 self.assertEqual(dataset_postcached[0]["label"], os.path.join(tempdir, "test_label1_new.nii.gz"))
                 self.assertEqual(dataset_postcached[1]["extra"], os.path.join(tempdir, "test_extra2_new.nii.gz"))
 
+    def test_different_transforms(self):
+        """
+        Different instances of `PersistentDataset` with the same cache_dir,
+        same input data, but different transforms should give different results.
+        """
+        shape = (1, 10, 9, 8)
+        im = np.arange(0, np.prod(shape)).reshape(shape)
+        with tempfile.TemporaryDirectory() as path:
+            im1 = PersistentDataset([im], Identity(), cache_dir=path, hash_transform=json_hashing)[0]
+            im2 = PersistentDataset([im], Flip(1), cache_dir=path, hash_transform=json_hashing)[0]
+            l2 = ((im1 - im2) ** 2).sum() ** 0.5
+            self.assertTrue(l2 > 1)
+
 
 if __name__ == "__main__":
     unittest.main()