5308 Add runtime cache support to CacheDataset

monai/apps/datasets.py

@@ -59,7 +59,7 @@ class MedNISTDataset(Randomizable, CacheDataset):
             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_workers: the number of worker threads to use.
+        num_workers: the number of worker threads if computing cache in the initialization.
             If num_workers is None then the number returned by os.cpu_count() is used.
             If a value less than 1 is specified, 1 will be used instead.
         progress: whether to display a progress bar when downloading dataset and computing the transform cache content.
@@ -70,6 +70,8 @@ class MedNISTDataset(Randomizable, CacheDataset):
             may set `copy=False` for better performance.
         as_contiguous: whether to convert the cached NumPy array or PyTorch tensor to be contiguous.
             it may help improve the performance of following logic.
+        runtime_cache: whether to compute cache at the runtime, default to `False` to prepare
+            the cache content at initializaiton.
 
     Raises:
         ValueError: When ``root_dir`` is not a directory.
@@ -97,6 +99,7 @@ def __init__(
         progress: bool = True,
         copy_cache: bool = True,
         as_contiguous: bool = True,
+        runtime_cache: bool = False,
     ) -> None:
         root_dir = Path(root_dir)
         if not root_dir.is_dir():
@@ -135,6 +138,7 @@ def __init__(
             progress=progress,
             copy_cache=copy_cache,
             as_contiguous=as_contiguous,
+            runtime_cache=runtime_cache,
         )
 
     def randomize(self, data: np.ndarray) -> None:
@@ -212,7 +216,7 @@ class DecathlonDataset(Randomizable, CacheDataset):
             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_workers: the number of worker threads to use.
+        num_workers: the number of worker threads if computing cache in the initialization.
             If num_workers is None then the number returned by os.cpu_count() is used.
             If a value less than 1 is specified, 1 will be used instead.
         progress: whether to display a progress bar when downloading dataset and computing the transform cache content.
@@ -223,6 +227,8 @@ class DecathlonDataset(Randomizable, CacheDataset):
             may set `copy=False` for better performance.
         as_contiguous: whether to convert the cached NumPy array or PyTorch tensor to be contiguous.
             it may help improve the performance of following logic.
+        runtime_cache: whether to compute cache at the runtime, default to `False` to prepare
+            the cache content at initializaiton.
 
     Raises:
         ValueError: When ``root_dir`` is not a directory.
@@ -290,6 +296,7 @@ def __init__(
         progress: bool = True,
         copy_cache: bool = True,
         as_contiguous: bool = True,
+        runtime_cache: bool = False,
     ) -> None:
         root_dir = Path(root_dir)
         if not root_dir.is_dir():
@@ -342,6 +349,7 @@ def __init__(
             progress=progress,
             copy_cache=copy_cache,
             as_contiguous=as_contiguous,
+            runtime_cache=runtime_cache,
         )
 
     def get_indices(self) -> np.ndarray:
@@ -438,7 +446,7 @@ class TciaDataset(Randomizable, CacheDataset):
             will take the minimum of (cache_num, data_length x cache_rate, data_length).
         cache_rate: percentage of cached data in total, default is 0.0 (no cache).
             will take the minimum of (cache_num, data_length x cache_rate, data_length).
-        num_workers: the number of worker threads to use.
+        num_workers: the number of worker threads if computing cache in the initialization.
             If num_workers is None then the number returned by os.cpu_count() is used.
             If a value less than 1 is specified, 1 will be used instead.
         progress: whether to display a progress bar when downloading dataset and computing the transform cache content.
@@ -449,6 +457,8 @@ class TciaDataset(Randomizable, CacheDataset):
             may set `copy=False` for better performance.
         as_contiguous: whether to convert the cached NumPy array or PyTorch tensor to be contiguous.
             it may help improve the performance of following logic.
+        runtime_cache: whether to compute cache at the runtime, default to `False` to prepare
+            the cache content at initializaiton.
 
     Example::
 
@@ -504,6 +514,7 @@ def __init__(
         progress: bool = True,
         copy_cache: bool = True,
         as_contiguous: bool = True,
+        runtime_cache: bool = False,
     ) -> None:
         root_dir = Path(root_dir)
         if not root_dir.is_dir():
@@ -550,6 +561,7 @@ def __init__(
             progress=progress,
             copy_cache=copy_cache,
             as_contiguous=as_contiguous,
+            runtime_cache=runtime_cache,
         )
 
     def get_indices(self) -> np.ndarray:

monai/data/dataloader.py

@@ -80,6 +80,12 @@ def __init__(self, dataset: Dataset, num_workers: int = 0, **kwargs) -> None:
             init_seed = _g.initial_seed()
             _seed = torch.empty((), dtype=torch.int64).random_(generator=_g).item()
             set_rnd(dataset, int(_seed))
+            # disable unnecessary multiprocessing caching
+            from monai.data.dataset import CacheDataset  # avoid circular import
+
+            if isinstance(dataset, CacheDataset) and dataset.runtime_cache:
+                dataset.disable_share_memory_cache()
+
             _g.manual_seed(init_seed)
         if "collate_fn" not in kwargs:
             kwargs["collate_fn"] = list_data_collate

monai/data/dataset.py

@@ -19,12 +19,15 @@
 import time
 import warnings
 from copy import copy, deepcopy
+from multiprocessing.managers import ListProxy
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
 from typing import IO, TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union
 
 import numpy as np
 import torch
+import torch.distributed as dist
+from torch.multiprocessing import Manager
 from torch.serialization import DEFAULT_PROTOCOL
 from torch.utils.data import Dataset as _TorchDataset
 from torch.utils.data import Subset
@@ -748,6 +751,7 @@ def __init__(
         as_contiguous: bool = True,
         hash_as_key: bool = False,
         hash_func: Callable[..., bytes] = pickle_hashing,
+        runtime_cache: bool = False,
     ) -> None:
         """
         Args:
@@ -757,7 +761,7 @@ def __init__(
                 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_workers: the number of worker threads to use.
+            num_workers: the number of worker threads if computing cache in the initialization.
                 If num_workers is None then the number returned by os.cpu_count() is used.
                 If a value less than 1 is speficied, 1 will be used instead.
             progress: whether to display a progress bar.
@@ -773,6 +777,14 @@ def __init__(
                 the dataset has duplicated items or augmented dataset.
             hash_func: if `hash_as_key`, a callable to compute hash from data items to be cached.
                 defaults to `monai.data.utils.pickle_hashing`.
+            runtime_cache: whether to compute cache at the runtime, default to `False` to prepare
+                the cache content at initializaiton, if `True`, it will cache during the first epoch
+                of model training, so it can start the first mini-batch earlier. please note that:
+                1. when using this option in multi-gpu distributed training,
+                `torch.cuda.set_device()` must be called before initializing this class.
+                2. to execute `runtime cache` on GPU memory, must co-work with
+                `monai.data.DataLoader`, and can't work with `monai.data.DistributedSampler`
+                as GPU Tensor usually can't be shared in the multiprocessing context.
 
         """
         if not isinstance(transform, Compose):
@@ -788,9 +800,11 @@ def __init__(
         self.num_workers = num_workers
         if self.num_workers is not None:
             self.num_workers = max(int(self.num_workers), 1)
+        self.runtime_cache = runtime_cache
         self.cache_num = 0
-        self._cache: List = []
+        self._cache: Union[List, ListProxy] = []
         self._hash_keys: List = []
+        self._is_dist = dist.is_available() and dist.is_initialized()
         self.set_data(data)
 
     def set_data(self, data: Sequence):
@@ -807,6 +821,18 @@ def set_data(self, data: Sequence):
         def _compute_cache_num(data_len: int):
             self.cache_num = min(int(self.set_num), int(data_len * self.set_rate), data_len)
 
+        def _compute_cache(indices=None):
+            if self.runtime_cache:
+                cache = Manager().list([None for _ in range(self.cache_num)])
+                if self._is_dist:
+                    obj_list = [cache]
+                    # broadcast the ProxyList to all the ranks, then share the same cache content at runtime
+                    dist.broadcast_object_list(obj_list, src=0)
+                    cache = obj_list[0]
+            else:
+                cache = self._fill_cache(indices)
+            return cache
+
         if self.hash_as_key:
             # only compute cache for the unique items of dataset, and record the last index for duplicated items
             mapping = {self.hash_func(v): i for i, v in enumerate(data)}
@@ -816,7 +842,16 @@ def _compute_cache_num(data_len: int):
         else:
             _compute_cache_num(len(self.data))
             indices = list(range(self.cache_num))
-        self._cache = self._fill_cache(indices)
+
+        self._cache = _compute_cache(indices)
+
+    def disable_share_memory_cache(self):
+        """
+        If the cache content is multiprocessing share memory list, convert it to a regular ptython list.
+        Because multiprocessing ProxyList is not supported for the GPU caching, may need to explicitly diasble it.
+
+        """
+        self._cache = list(self._cache)
 
     def _fill_cache(self, indices=None) -> List:
         """
@@ -871,6 +906,10 @@ def _transform(self, index: int):
         if self._cache is None:
             raise RuntimeError("cache buffer is not initialized, please call `set_data()` first.")
         data = self._cache[cache_index]
+        # runtime cache computation
+        if data is None:
+            data = self._cache[cache_index] = self._load_cache_item(cache_index)
+
         # load data from cache and execute from the first random transform
         start_run = False
         if not isinstance(self.transform, Compose):
@@ -973,6 +1012,7 @@ def __init__(
         seed: int = 0,
         copy_cache: bool = True,
         as_contiguous: bool = True,
+        runtime_cache: bool = False,
     ) -> None:
         if shuffle:
             self.set_random_state(seed=seed)

tests/test_decathlondataset.py

@@ -47,7 +47,12 @@ def _test_dataset(dataset):
 
         _test_dataset(data)
         data = DecathlonDataset(
-            root_dir=testing_dir, task="Task04_Hippocampus", transform=transform, section="validation", download=False
+            root_dir=testing_dir,
+            task="Task04_Hippocampus",
+            transform=transform,
+            section="validation",
+            download=False,
+            runtime_cache=True,
         )
         _test_dataset(data)
         self.assertTrue(data[0]["image"].meta["filename_or_obj"].endswith("hippocampus_163.nii.gz"))

tests/test_integration_fast_train.py

@@ -144,7 +144,7 @@ def test_train_timing(self):
 
         # set CacheDataset, ThreadDataLoader and DiceCE loss for MONAI fast training
         train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=8)
-        val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=5)
+        val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, runtime_cache=True)
         # disable multi-workers because `ThreadDataLoader` works with multi-threads
         train_loader = ThreadDataLoader(train_ds, num_workers=0, batch_size=4, shuffle=True)
         val_loader = ThreadDataLoader(val_ds, num_workers=0, batch_size=1)

tests/test_integration_segmentation_3d.py

@@ -83,7 +83,9 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None,
 
     # create a training data loader
     if cachedataset == 2:
-        train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.8)
+        train_ds = monai.data.CacheDataset(
+            data=train_files, transform=train_transforms, cache_rate=0.8, runtime_cache=True
+        )
     elif cachedataset == 3:
         train_ds = monai.data.LMDBDataset(data=train_files, transform=train_transforms, cache_dir=root_dir)
     else:
@@ -170,7 +172,7 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None,
                 plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="output")
     print(f"train completed, best_metric: {best_metric:0.4f}  at epoch: {best_metric_epoch}")
     writer.close()
-    return epoch_loss_values, best_metric, best_metric_epoch
+    return epoch_loss_values, best_metric
 
 
 def run_inference_test(root_dir, device="cuda:0"):
@@ -256,9 +258,7 @@ def train_and_infer(self, idx=0):
             _readers = ("itkreader", "itkreader")
         elif idx == 2:
             _readers = ("itkreader", "nibabelreader")
-        losses, best_metric, best_metric_epoch = run_training_test(
-            self.data_dir, device=self.device, cachedataset=idx, readers=_readers
-        )
+        losses, best_metric = run_training_test(self.data_dir, device=self.device, cachedataset=idx, readers=_readers)
         infer_metric = run_inference_test(self.data_dir, device=self.device)
 
         # check training properties

tests/test_mednistdataset.py

@@ -43,7 +43,9 @@ def _test_dataset(dataset):
         _test_dataset(data)
 
         # testing from
-        data = MedNISTDataset(root_dir=Path(testing_dir), transform=transform, section="test", download=False)
+        data = MedNISTDataset(
+            root_dir=Path(testing_dir), transform=transform, section="test", download=False, runtime_cache=True
+        )
         self.assertEqual(data.get_num_classes(), 6)
         _test_dataset(data)
         data = MedNISTDataset(root_dir=testing_dir, section="test", download=False)

tests/test_sampler_dist.py

@@ -12,9 +12,11 @@
 import unittest
 
 import numpy as np
+import torch
 import torch.distributed as dist
 
-from monai.data import DistributedSampler
+from monai.data import CacheDataset, DataLoader, DistributedSampler
+from monai.transforms import ToTensor
 from tests.utils import DistCall, DistTestCase
 
 
@@ -43,6 +45,27 @@ def test_uneven(self):
         if dist.get_rank() == 1:
             np.testing.assert_allclose(samples, np.array([2, 4]))
 
+    @DistCall(nnodes=1, nproc_per_node=2)
+    def test_cachedataset(self):
+        data = [1, 2, 3, 4, 5]
+        dataset = CacheDataset(data=data, transform=ToTensor(track_meta=False), cache_rate=1.0, runtime_cache=True)
+        sampler = DistributedSampler(dataset=dataset, shuffle=False, even_divisible=False)
+        dataloader = DataLoader(dataset=dataset, sampler=sampler, batch_size=1, num_workers=2)
+        for i in range(3):
+            dist.barrier()
+            if i > 0:
+                # verify the runtime cache content is completed after first epoch
+                for j, c in enumerate(dataset._cache):
+                    self.assertTrue(isinstance(c, torch.Tensor))
+                    torch.testing.assert_allclose(c, j + 1)
+            for k, d in enumerate(dataloader):
+                self.assertTrue(isinstance(d, torch.Tensor))
+                if dist.get_rank() == 0:
+                    torch.testing.assert_allclose(d[0], k * 2 + 1)
+
+                if dist.get_rank() == 1:
+                    torch.testing.assert_allclose(d[0], (k + 1) * 2)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_tciadataset.py

@@ -64,6 +64,7 @@ def _test_dataset(dataset):
             section="validation",
             download=False,
             val_frac=val_frac,
+            runtime_cache=True,
         )
         _test_dataset(data)
         self.assertTrue(