diff --git a/monai/apps/datasets.py b/monai/apps/datasets.py index e5992711da..cd2537a4c6 100644 --- a/monai/apps/datasets.py +++ b/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: diff --git a/monai/data/dataloader.py b/monai/data/dataloader.py index f43211f184..53670046d8 100644 --- a/monai/data/dataloader.py +++ b/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 diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 2c263b3e32..00b113eda2 100644 --- a/monai/data/dataset.py +++ b/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) diff --git a/tests/test_decathlondataset.py b/tests/test_decathlondataset.py index 49280f6fa6..b3844484ce 100644 --- a/tests/test_decathlondataset.py +++ b/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")) diff --git a/tests/test_integration_fast_train.py b/tests/test_integration_fast_train.py index 13f918d201..bb50ddf7b6 100644 --- a/tests/test_integration_fast_train.py +++ b/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) diff --git a/tests/test_integration_segmentation_3d.py b/tests/test_integration_segmentation_3d.py index b5b1d69565..7505bda03c 100644 --- a/tests/test_integration_segmentation_3d.py +++ b/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 diff --git a/tests/test_mednistdataset.py b/tests/test_mednistdataset.py index 87a0d6a2ec..c46d846903 100644 --- a/tests/test_mednistdataset.py +++ b/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) diff --git a/tests/test_sampler_dist.py b/tests/test_sampler_dist.py index 8b140f3ff8..1608754da3 100644 --- a/tests/test_sampler_dist.py +++ b/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() diff --git a/tests/test_tciadataset.py b/tests/test_tciadataset.py index 4da897f14b..301f2caac2 100644 --- a/tests/test_tciadataset.py +++ b/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(