Repository: Project-MONAI/MONAI

Issue:
Use persistent caching for pre-processing
**Feature request**
Pre-processing large data sets (10000 3D medical images) can be a significant time burden and is too large to fit completely in memory.

**Describe the solution you'd like**
A persistent cache of pre-processed data that is generated on-the-fly as needed so that initial testing can occur quickly.

**Describe alternatives you've considered**
In-memory caching alone is insufficient for meeting the efficiency needs of experimentation of parameter spaces or simultaneously running many model configurations on the same pre-processed data.

Investigated Dask, and Partd, and similar parallel data stores, but those solutions seemed more invasive and less compatible with the current infrastructure.

**Additional context**
I originally was running out of memory on my 756GB memory node when attempting to cache results. 


Discussion:
Thanks very much for your valuable suggestion, it will certainly improve the cache capability if supporting persistent storage.
Let's investigate it when @gml16 completed the multi-thread cache.
Of course, if you have interest to submit a proposal or PR, do welcome!
Thanks.
@Nic-Ma PR will be arriving in a few minutes.  It will likely have merge conflicts with @gml16, but it should be easy to reconcile once each one or the other is approved.

See #317 

Resolving patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from .csv_saver import CSVSaver
-from .dataset import Dataset, CacheDataset
+from .dataset import Dataset, PersistentDataset, CacheDataset
 from .grid_dataset import GridPatchDataset
 from .nifti_reader import load_nifti, NiftiDataset
 from .nifti_saver import NiftiSaver
diff --git a/monai/data/dataset.py b/monai/data/dataset.py
--- a/monai/data/dataset.py
+++ b/monai/data/dataset.py
@@ -11,6 +11,10 @@
 
 import sys
 
+import hashlib
+import json
+from pathlib import Path
+
 import torch
 
 from multiprocessing.pool import ThreadPool
@@ -19,6 +23,7 @@
 from monai.transforms import Compose, Randomizable
 from monai.transforms.utils import apply_transform
 from monai.utils import process_bar
+from monai.utils.misc import ensure_tuple
 
 
 class Dataset(torch.utils.data.Dataset):
@@ -40,7 +45,10 @@ def __init__(self, data, transform=None):
             transform (Callable, optional): transforms to execute operations on input data.
         """
         self.data = data
-        self.transform = transform
+        if isinstance(transform, Compose):
+            self.transform = transform
+        else:
+            self.transform = Compose(ensure_tuple(transform))
 
     def __len__(self):
         return len(self.data)
@@ -53,6 +61,135 @@ def __getitem__(self, index):
         return data
 
 
+class PersistentDataset(Dataset):
+    """
+    Persistent storage of pre-computed values to efficiently manage larger than memory dictionary format data,
+    it can operate transforms for specific fields.  Results from the non-random transform components are computed
+    when first used, and stored in the `cache_dir` for rapid retrieval on subsequent uses.
+
+    For example, typical input data can be a list of dictionaries::
+
+        [{                            {                            {
+             'img': 'image1.nii.gz',      'img': 'image2.nii.gz',      'img': 'image3.nii.gz',
+             'seg': 'label1.nii.gz',      'seg': 'label2.nii.gz',      'seg': 'label3.nii.gz',
+             'extra': 123                 'extra': 456                 'extra': 789
+         },                           },                           }]
+
+    For a composite transform like 
+     [
+        LoadNiftid(keys=['image', 'label']),
+        Orientationd(keys=['image', 'label'], axcodes='RAS'),
+        ScaleIntensityRanged(keys=['image'], a_min=-57, a_max=164, b_min=0.0, b_max=1.0, clip=True),
+        RandCropByPosNegLabeld(keys=['image', 'label'], label_key='label', size=(96, 96, 96), pos=1,
+                           neg=1, num_samples=4, image_key='image', image_threshold=0),
+        ToTensord(keys=['image', 'label'])
+     ]
+
+     Upon first use a filename based dataset will be processed by the transform for the
+     [LoadNiftid, Orientationd, ScaleIntensityRanged] and the resulting tensor written to
+     the `cache_dir` before applying the remainging random dependant transforms
+     [RandCropByPosNegLabeld, ToTensord] elements for use in the analysis.
+
+     Subsequent uses of a dataset directly read pre-processed results from `cache_dir`
+     followed by applying the random dependant parts of transform processing.
+
+    """
+
+    def __init__(self, data, transform=None, cache_dir=None):
+        """
+        Args:
+            data (Iterable): input data to load and transform to generate dataset for model.
+            transform (Callable, optional): transforms to execute operations on input data.
+            cache_dir (Path or str or None): If specified, this is the location for persistent storage
+                of pre-computed transformed data tensors. The cache_dir is computed once, and
+                persists on disk until explicitly removed.  Different runs, programs, experiments
+                may share a common cache dir provided that the trasnsforms pre-processing is
+                consistent.
+        """
+        super().__init__(data=data, transform=transform)
+        self.cache_dir = Path(cache_dir) if cache_dir is not None else None
+
+    def _pre_first_random_transform(self, item_transformed):
+        """
+        Process the data from original state up to the first random element.
+
+        Args:
+            item_transformed: The data to be transformed
+        Returns:
+            the transformed element up to the first identified
+            random transform object
+        """
+        for _transform in self.transform.transforms:
+            # execute all the deterministic transforms before the first random transform
+            if isinstance(_transform, Randomizable):
+                break
+            item_transformed = apply_transform(_transform, item_transformed)
+        return item_transformed
+
+    def _first_random_and_beyond_transform(self, item_transformed):
+        """
+        Process the data from before the first random transform to the final state ready for evaluation.
+        Args:
+            item_transformed: The data to be transformed (already process upto the first random transform)
+        Returns:
+            the transformed element through the random transforms
+        """
+        start_post_randomize_run = False
+        for _transform in self.transform.transforms:
+            if start_post_randomize_run or isinstance(_transform, Randomizable):
+                start_post_randomize_run = True
+                item_transformed = apply_transform(_transform, item_transformed)
+        return item_transformed
+
+    def _pre_first_random_cachecheck(self, item_transformed):
+        """
+            A function to cache the expensive input data transform operations
+            so that huge data sets (larger than computer memory) can be processed
+            on the fly as needed, and intermediate results written to disk for
+            future use.
+        Args:
+            item_transformed: The current data element to be mutated into transformed representation
+
+        Returns:
+            The transformed data_element, either from cache, or explicitly computing it.
+
+        Warning:
+            The current implementation does not encode transform information as part of the
+            hashing mechanism used for generating cache names.  If the transforms applied are
+            changed in any way, the objects in the cache dir will be invalid.  The hash for the
+            cache is ONLY dependant on the input filename paths.
+        """
+        if item_transformed.get("cached", False) is False:
+            hashfile = None
+            if self.cache_dir is not None:
+                cache_dir_path: Path = Path(self.cache_dir)
+                if cache_dir_path.is_dir():
+                    # TODO: Find way to hash transforms content as part of the cache
+                    data_item_md5 = hashlib.md5(json.dumps(item_transformed, sort_keys=True).encode('utf-8')).hexdigest()
+                    hashfile: Path = Path(cache_dir_path) / f"{data_item_md5}.pt"
+
+            if hashfile is not None and hashfile.is_file():
+                item_transformed = torch.load(hashfile)
+            else:
+                item_transformed = self._pre_first_random_transform(item_transformed)
+                if hashfile is not None:
+                    # add sentinal flag to indicate that the transforms have already been computed.
+                    item_transformed["cache"] = True
+                    # NOTE: Writing to ".temp_write_cache" and then using a nearly atomic rename operation
+                    #       to make the cache more robust to manual killing of parent process
+                    #       which may leave partially written cache files in an incomplete state
+                    temp_hash_file: Path = hashfile.with_suffix(".temp_write_cache")
+                    torch.save(item_transformed, temp_hash_file)
+                    temp_hash_file.rename(hashfile)
+
+        return item_transformed
+
+    def __getitem__(self, index):
+        pre_random_item = self._pre_first_random_cachecheck(self.data[index])
+        post_random_item = self._first_random_and_beyond_transform(pre_random_item)
+        return post_random_item
+
+
 class CacheDataset(Dataset):
     """
     Dataset with cache mechanism that can load data and cache deterministic transforms' result during training.


Test patch:
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -11,7 +11,6 @@
 
 import unittest
 import os
-import shutil
 import numpy as np
 import tempfile
 import nibabel as nib
@@ -23,7 +22,6 @@
     (128, 128, 128)
 ]
 
-
 class TestDataset(unittest.TestCase):
 
     @parameterized.expand([TEST_CASE_1])
@@ -54,7 +52,7 @@ def test_shape(self, expected_shape):
         dataset = Dataset(data=test_data, transform=test_transform)
         data1 = dataset[0]
         data2 = dataset[1]
-        shutil.rmtree(tempdir)
+
         self.assertTupleEqual(data1['image'].shape, expected_shape)
         self.assertTupleEqual(data1['label'].shape, expected_shape)
         self.assertTupleEqual(data1['extra'].shape, expected_shape)
@@ -62,6 +60,17 @@ def test_shape(self, expected_shape):
         self.assertTupleEqual(data2['label'].shape, expected_shape)
         self.assertTupleEqual(data2['extra'].shape, expected_shape)
 
+        dataset = Dataset(data=test_data, transform=LoadNiftid(keys=['image', 'label', 'extra']))
+        data1_simple = dataset[0]
+        data2_simple = dataset[1]
+
+        self.assertTupleEqual(data1_simple['image'].shape, expected_shape)
+        self.assertTupleEqual(data1_simple['label'].shape, expected_shape)
+        self.assertTupleEqual(data1_simple['extra'].shape, expected_shape)
+        self.assertTupleEqual(data2_simple['image'].shape, expected_shape)
+        self.assertTupleEqual(data2_simple['label'].shape, expected_shape)
+        self.assertTupleEqual(data2_simple['extra'].shape, expected_shape)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_persistentdataset.py b/tests/test_persistentdataset.py
new file mode 100644
--- /dev/null
+++ b/tests/test_persistentdataset.py
@@ -0,0 +1,80 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import shutil
+import numpy as np
+import tempfile
+import nibabel as nib
+from parameterized import parameterized
+from monai.data import PersistentDataset
+from monai.transforms import Compose, LoadNiftid, SimulateDelayd
+
+TEST_CASE_1 = [
+    (128, 128, 128)
+]
+
+
+class TestDataset(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1])
+    def test_shape(self, expected_shape):
+        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
+        tempdir = tempfile.mkdtemp()
+        nib.save(test_image, os.path.join(tempdir, 'test_image1.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_label1.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_extra1.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_image2.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_label2.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_extra2.nii.gz'))
+        test_data = [
+            {
+                'image': os.path.join(tempdir, 'test_image1.nii.gz'),
+                'label': os.path.join(tempdir, 'test_label1.nii.gz'),
+                'extra': os.path.join(tempdir, 'test_extra1.nii.gz')
+            },
+            {
+                'image': os.path.join(tempdir, 'test_image2.nii.gz'),
+                'label': os.path.join(tempdir, 'test_label2.nii.gz'),
+                'extra': os.path.join(tempdir, 'test_extra2.nii.gz')
+            }
+        ]
+        test_transform = Compose([LoadNiftid(keys=['image', 'label', 'extra']),
+                                  SimulateDelayd(keys=['image', 'label', 'extra'],
+                                                 delay_time=[1e-7, 1e-6, 1e-5])])
+
+        dataset_precached = PersistentDataset(data=test_data, transform=test_transform, cache_dir=tempdir)
+        data1_precached = dataset_precached[0]
+        data2_precached = dataset_precached[1]
+
+        dataset_postcached = PersistentDataset(data=test_data, transform=test_transform, cache_dir=tempdir)
+        data1_postcached = dataset_postcached[0]
+        data2_postcached = dataset_postcached[1]
+        shutil.rmtree(tempdir)
+
+        self.assertTupleEqual(data1_precached['image'].shape, expected_shape)
+        self.assertTupleEqual(data1_precached['label'].shape, expected_shape)
+        self.assertTupleEqual(data1_precached['extra'].shape, expected_shape)
+        self.assertTupleEqual(data2_precached['image'].shape, expected_shape)
+        self.assertTupleEqual(data2_precached['label'].shape, expected_shape)
+        self.assertTupleEqual(data2_precached['extra'].shape, expected_shape)
+
+        self.assertTupleEqual(data1_postcached['image'].shape, expected_shape)
+        self.assertTupleEqual(data1_postcached['label'].shape, expected_shape)
+        self.assertTupleEqual(data1_postcached['extra'].shape, expected_shape)
+        self.assertTupleEqual(data2_postcached['image'].shape, expected_shape)
+        self.assertTupleEqual(data2_postcached['label'].shape, expected_shape)
+        self.assertTupleEqual(data2_postcached['extra'].shape, expected_shape)
+
+
+if __name__ == '__main__':
+    unittest.main()