Repository: Project-MONAI/MONAI

Issue:
Develop SmartCacheDataset based on CacheDataset
**Is your feature request related to a problem? Please describe.**
For very big dataset, we should support to only cache and use part of the dataset to train for every epoch, create a background thread/process to replace the cached content for every epoch.
Need to handle the multi-threads and multi-processing issues clearly in this task.

see also https://docs.nvidia.com/clara/tlt-mi/clara-train-sdk-v3.0/nvmidl/additional_features/smart_cache.html#smart-cache


Discussion:

Resolving patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -11,7 +11,7 @@
 
 from .csv_saver import CSVSaver
 from .dataloader import DataLoader
-from .dataset import ArrayDataset, CacheDataset, Dataset, PersistentDataset, ZipDataset
+from .dataset import ArrayDataset, CacheDataset, Dataset, PersistentDataset, SmartCacheDataset, ZipDataset
 from .decathalon_datalist import load_decathalon_datalist
 from .grid_dataset import GridPatchDataset
 from .image_reader import *
diff --git a/monai/data/dataset.py b/monai/data/dataset.py
--- a/monai/data/dataset.py
+++ b/monai/data/dataset.py
@@ -11,11 +11,13 @@
 
 import hashlib
 import json
+import math
 import sys
 import threading
+import time
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
-from typing import Any, Callable, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -261,7 +263,7 @@ def __init__(
         if not isinstance(transform, Compose):
             transform = Compose(transform)
         super().__init__(data, transform)
-        self.cache_num = min(cache_num, int(len(self) * cache_rate), len(self))
+        self.cache_num = min(cache_num, int(len(data) * cache_rate), len(data))
         if self.cache_num > 0:
             self._cache = [None] * self.cache_num
             if num_workers > 0:
@@ -321,6 +323,238 @@ def __getitem__(self, index):
         return data
 
 
+class SmartCacheDataset(CacheDataset):
+    """
+    Re-implementation of the SmartCache mechanism in NVIDIA Clara-train SDK.
+    At any time, the cache pool only keeps a subset of the whole dataset. In each epoch, only the items
+    in the cache are used for training. This ensures that data needed for training is readily available,
+    keeping GPU resources busy. Note that cached items may still have to go through a non-deterministic
+    transform sequence before being fed to GPU. At the same time, another thread is preparing replacement
+    items by applying the transform sequence to items not in cache. Once one epoch is completed, Smart
+    Cache replaces the same number of items with replacement items.
+    Smart Cache uses a simple `running window` algorithm to determine the cache content and replacement items.
+    Let N be the configured number of objects in cache; and R be the number of replacement objects (R = ceil(N * r),
+    where r is the configured replace rate).
+    For more details, please refer to:
+    https://docs.nvidia.com/clara/tlt-mi/clara-train-sdk-v3.0/nvmidl/additional_features/smart_cache.html#smart-cache
+
+    For example, if we have 5 images: `[image1, image2, image3, image4, image5]`, and `cache_num=4`, `replace_rate=0.25`.
+    so the actual training images cached and replaced for every epoch are as below:
+    epoch 1: [image1, image2, image3, image4]
+    epoch 2: [image2, image3, image4, image5]
+    epoch 3: [image3, image4, image5, image1]
+    epoch 3: [image4, image5, image1, image2]
+    epoch N: [image[N % 5] ...]
+
+    The usage of `SmartCacheDataset` contains 4 steps:
+    1. Initialize `SmartCacheDataset` object and cache for the first epoch.
+    2. Call `start()` to run replacement thread in background.
+    3. Call `update_cache()` before every epoch to replace training items.
+    4. Call `shutdown()` when training ends.
+
+    Args:
+        data: input data to load and transform to generate dataset for model.
+        transform: transforms to execute operations on input data.
+        replace_rate: percentage of the cached items to be replaced in every epoch.
+        cache_num: number of items to be cached. Default is `sys.maxsize`.
+            will take the minimum of (cache_num, data_length x cache_rate, data_length).
+        cache_rate: percentage of cached data in total, default is 1.0 (cache all).
+            will take the minimum of (cache_num, data_length x cache_rate, data_length).
+        num_init_workers: the number of worker threads to initialize the cache for first epoch.
+            if 0, run in main thread, no separate thread will open.
+        num_replace_workers: the number of worker threads to prepare the replacement cache for every epoch.
+            if 0, run in main thread, no separate thread will open.
+
+    """
+
+    def __init__(
+        self,
+        data: Sequence,
+        transform: Union[Sequence[Callable], Callable],
+        replace_rate: float,
+        cache_num: int = sys.maxsize,
+        cache_rate: float = 1.0,
+        num_init_workers: int = 0,
+        num_replace_workers: int = 0,
+    ) -> None:
+        super().__init__(data, transform, cache_num, cache_rate, num_init_workers)
+        if self.cache_num == len(data):
+            raise ValueError("cache_num must be smaller than dataset length to support replacement.")
+        if replace_rate <= 0:
+            raise ValueError("replace_rate must be greater than 0, otherwise, please use CacheDataset.")
+        self.num_replace_workers: int = num_replace_workers
+
+        self._total_num: int = len(data)
+        self._replace_num: int = min(math.ceil(self.cache_num * replace_rate), len(data) - self.cache_num)
+        self._replacements: List[Any] = [None for _ in range(self._replace_num)]
+        self._replace_data_idx: List[int] = list(range(self._replace_num))
+
+        self._start_pos: int = 0
+        self._update_lock: threading.Lock = threading.Lock()
+        self._round: int = 1
+        self._replace_done: bool = False
+        self._replace_mgr: Optional[threading.Thread] = None
+
+        self._compute_data_idx()
+
+    def _compute_data_idx(self):
+        """
+        Update the replacement data position in the total data.
+
+        """
+        for i in range(self._replace_num):
+            pos: int = self._start_pos + self.cache_num + i
+            if pos >= self._total_num:
+                pos -= self._total_num
+            self._replace_data_idx[i] = pos
+
+    def is_started(self):
+        """
+        Check whether the replacement thread is already started.
+
+        """
+        if self._replace_mgr is None:
+            return False
+        return self._replace_mgr.isAlive()
+
+    def start(self):
+        """
+        Start the background thread to replace training items for every epoch.
+
+        """
+        if self._replace_mgr is None or not self._is_started():
+            self._restart()
+
+    def _restart(self):
+        """
+        Restart background thread if killed for some reason.
+
+        """
+        self._round = 1
+        self._replace_mgr = threading.Thread(target=self.manage_replacement, daemon=True)
+        self._replace_mgr.start()
+
+    def _try_update_cache(self):
+        """
+        Update the cache items with new replacement for current epoch.
+
+        """
+        with self._update_lock:
+            if self._replace_done:
+                remain_num: int = self.cache_num - self._replace_num
+                for i in range(remain_num):
+                    self._cache[i] = self._cache[i + self._replace_num]
+                for i in range(self._replace_num):
+                    self._cache[remain_num + i] = self._replacements[i]
+
+                self._start_pos += self._replace_num
+                if self._start_pos >= self._total_num:
+                    self._start_pos -= self._total_num
+
+                self._compute_data_idx()
+
+                # ready for next round
+                self._round += 1
+                self._replace_done = False
+                return True
+            else:
+                return False
+
+    def update_cache(self):
+        """
+        Update cache items for current epoch, need to call this function before every epoch.
+        If the cache has been shutdown before, need to restart the `_replace_mgr` thread.
+
+        """
+        if not self._replace_mgr.isAlive():
+            self._restart()
+
+        # make sure update is done
+        while not self._try_update_cache():
+            time.sleep(0.01)
+
+    def _try_shutdown(self):
+        """
+        Wait for thread lock to shut down the background thread.
+
+        """
+        with self._update_lock:
+            if self._replace_done:
+                self._round = 0
+                self._replace_done = False
+                return True
+            else:
+                return False
+
+    def shutdown(self):
+        """
+        Shut down the background thread for replacement.
+
+        """
+        if not self.is_started():
+            return
+
+        # wait until replace mgr is done the current round
+        while not self._try_shutdown():
+            time.sleep(0.01)
+        self._replace_mgr.join()
+
+    def _replace_cache_thread(self, index: int):
+        """
+        Execute deterministic transforms on the new data for replacement.
+
+        """
+        pos: int = self._replace_data_idx[index]
+        self._replacements[index] = self._load_cache_item(self.data[pos], self.transform.transforms)  # type: ignore
+
+    def _compute_replacements(self):
+        """
+        Compute expected items for the replacement of next epoch, execute deterministic transforms.
+        It can support multi-threads to accelerate the computation progress.
+
+        """
+        if self.num_replace_workers > 0:
+            with ThreadPool(self.num_replace_workers) as p:
+                p.map(self._replace_cache_thread, list(range(self._replace_num)))
+        else:
+            for i in range(self._replace_num):
+                self._replace_cache_thread(i)
+        self._replace_done = True
+
+    def _try_manage_replacement(self, check_round):
+        """
+        Wait thread lock and replace training items in the background thread.
+
+        """
+        with self._update_lock:
+            if self._round <= 0:
+                # shutdown replacement
+                self._replace_done = True
+                return True, -1
+
+            if self._round != check_round:
+                self._compute_replacements()
+            return False, self._round
+
+    def manage_replacement(self):
+        """
+        Background thread for replacement.
+
+        """
+        check_round: int = -1
+        done = False
+        while not done:
+            done, check_round = self._try_manage_replacement(check_round)
+            time.sleep(0.01)
+
+    def __len__(self):
+        """
+        The dataset length is given by cache_num instead of len(data).
+
+        """
+        return self.cache_num
+
+
 class ZipDataset(Dataset):
     """
     Zip several PyTorch datasets and output data(with the same index) together in a tuple.


Test patch:
diff --git a/tests/test_smartcachedataset.py b/tests/test_smartcachedataset.py
new file mode 100644
--- /dev/null
+++ b/tests/test_smartcachedataset.py
@@ -0,0 +1,72 @@
+# 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 os
+import tempfile
+import unittest
+
+import nibabel as nib
+import numpy as np
+from parameterized import parameterized
+
+from monai.data import SmartCacheDataset
+from monai.transforms import Compose, LoadNiftid
+
+TEST_CASE_1 = [0.1, 0, Compose([LoadNiftid(keys=["image", "label", "extra"])])]
+
+TEST_CASE_2 = [0.1, 4, Compose([LoadNiftid(keys=["image", "label", "extra"])])]
+
+TEST_CASE_3 = [0.1, 4, None]
+
+TEST_CASE_4 = [0.5, 2, Compose([LoadNiftid(keys=["image", "label", "extra"])])]
+
+
+class TestSmartCacheDataset(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    def test_shape(self, replace_rate, num_replace_workers, transform):
+        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), 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"))
+            nib.save(test_image, os.path.join(tempdir, "test_extra1.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"),
+                }
+            ] * 20
+            dataset = SmartCacheDataset(
+                data=test_data,
+                transform=transform,
+                replace_rate=replace_rate,
+                cache_num=16,
+                num_init_workers=4,
+                num_replace_workers=num_replace_workers,
+            )
+
+            self.assertEqual(len(dataset._cache), dataset.cache_num)
+            for i in range(dataset.cache_num):
+                self.assertIsNotNone(dataset._cache[i])
+
+            dataset.start()
+            for _ in range(3):
+                dataset.update_cache()
+                self.assertIsNotNone(dataset._cache[15])
+                if isinstance(dataset._cache[15]["image"], np.ndarray):
+                    np.testing.assert_allclose(dataset._cache[15]["image"], dataset._cache[15]["label"])
+                else:
+                    self.assertIsInstance(dataset._cache[15]["image"], str)
+            dataset.shutdown()
+
+
+if __name__ == "__main__":
+    unittest.main()