Repository: Project-MONAI/MONAI

Issue:
Lightning Segmentation tutorial breaks on multiple GPU's (CacheDataset?)
Reproducing the spleen Lightning segmentation tutorial from [here](https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/spleen_segmentation_3d_lightning.ipynb). Full logs of two errors (1) and (2) shared below. You can reproduce with full code from the notebook into a singly python file can be copied from (3)
With the latest library versions, and running on two GPU's results in the following error (1):
In multiprocessing/reduction.py > def dump:
![image](https://user-images.githubusercontent.com/10415569/96673444-95c9d300-1377-11eb-83c9-ddc161190f27.png)
After reading full log (1) I've commented out `CropForegroundd(keys=["image", "label"], source_key="image"),` from both `train_transforms` and `val_transforms` and I get the following error (2):
In multiprocessing/reduction.py > def dump:
![image](https://user-images.githubusercontent.com/10415569/96673846-8303ce00-1378-11eb-97e8-1956d7208782.png)



(1) Full log:

> Exception has occurred: PicklingError
> Can't pickle <function CropForegroundd.<lambda> at 0x7fd6e5159b80>: attribute lookup CropForegroundd.<lambda> on monai.transforms.croppad.dictionary failed
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
>     ForkingPickler(file, protocol).dump(obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
>     reduction.dump(process_obj, fp)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
>     self._launch(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
>     super().__init__(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
>     return Popen(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/process.py", line 121, in start
>     self._popen = self._Popen(self)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 149, in start_processes
>     process.start()
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
>     return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_spawn_backend.py", line 64, in train
>     mp.spawn(self.ddp_train, nprocs=self.nprocs, args=(self.mp_queue, model,))
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
>     results = self.accelerator_backend.train()
>   File "/nfs/users/ext_bart.michiels/projects/examples/test_multi_gpu.py", line 213, in <module>
>     trainer.fit(net)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 97, in _run_module_code
>     _run_code(code, mod_globals, init_globals,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 265, in run_path
>     return _run_module_code(code, init_globals, run_name,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 194, in _run_module_as_main
>     return _run_code(code, main_globals, None,

(2) Full log:

> Exception has occurred: TypeError
> cannot pickle '_thread.lock' object
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
>     ForkingPickler(file, protocol).dump(obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
>     reduction.dump(process_obj, fp)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
>     self._launch(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
>     super().__init__(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
>     return Popen(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/process.py", line 121, in start
>     self._popen = self._Popen(self)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 149, in start_processes
>     process.start()
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
>     return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_spawn_backend.py", line 64, in train
>     mp.spawn(self.ddp_train, nprocs=self.nprocs, args=(self.mp_queue, model,))
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
>     results = self.accelerator_backend.train()
>   File "/nfs/users/ext_bart.michiels/projects/examples/test_multi_gpu.py", line 213, in <module>
>     trainer.fit(net)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 97, in _run_module_code
>     _run_code(code, mod_globals, init_globals,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 265, in run_path
>     return _run_module_code(code, init_globals, run_name,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 194, in _run_module_as_main
>     return _run_code(code, main_globals, None,

(3) Full python code from the notebook to reproduce the errors:

> import glob
> import os
> import shutil
> import tempfile
> import matplotlib.pyplot as plt
> import pytorch_lightning
> import torch
> from monai.apps import download_and_extract
> from monai.config import print_config
> from monai.data import CacheDataset, list_data_collate
> from monai.inferers import sliding_window_inference
> from monai.losses import DiceLoss
> from monai.metrics import compute_meandice
> from monai.networks.layers import Norm
> from monai.networks.nets import UNet
> from monai.transforms import (
>     AddChanneld,
>     Compose,
>     CropForegroundd,
>     LoadNiftid,
>     Orientationd,
>     RandCropByPosNegLabeld,
>     ScaleIntensityRanged,
>     Spacingd,
>     ToTensord,
> )
> from monai.utils import set_determinism
> print_config()
> directory = os.environ.get("MONAI_DATA_DIRECTORY")
> root_dir = tempfile.mkdtemp() if directory is None else directory
> print(root_dir)
> resource = "https://drive.google.com/uc?id=1jzeNU1EKnK81PyTsrx0ujfNl-t0Jo8uE"
> md5 = "410d4a301da4e5b2f6f86ec3ddba524e"
> compressed_file = os.path.join(root_dir, "Task09_Spleen.tar")
> data_dir = os.path.join(root_dir, "Task09_Spleen")
> if not os.path.exists(data_dir):
>     download_and_extract(resource, compressed_file, root_dir, md5)
> class Net(pytorch_lightning.LightningModule):
>     def __init__(self):
>         super().__init__()
>         self._model = UNet(
>             dimensions=3,
>             in_channels=1,
>             out_channels=2,
>             channels=(16, 32, 64, 128, 256),
>             strides=(2, 2, 2, 2),
>             num_res_units=2,
>             norm=Norm.BATCH,
>         )
>         self.loss_function = DiceLoss(to_onehot_y=True, softmax=True)
>         self.best_val_dice = 0
>         self.best_val_epoch = 0
> 
>     def forward(self, x):
>         return self._model(x)
> 
>     def prepare_data(self):
>         # set up the correct data path
>         train_images = sorted(glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
>         train_labels = sorted(glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
>         data_dicts = [
>             {"image": image_name, "label": label_name}
>             for image_name, label_name in zip(train_images, train_labels)
>         ]
>         train_files, val_files = data_dicts[:-9], data_dicts[-9:]
> 
>         # set deterministic training for reproducibility
>         set_determinism(seed=0)
> 
>         # define the data transforms
>         train_transforms = Compose(
>             [
>                 LoadNiftid(keys=["image", "label"]),
>                 AddChanneld(keys=["image", "label"]),
>                 Spacingd(
>                     keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest"),
>                 ),
>                 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,
>                 ),
>                 #CropForegroundd(keys=["image", "label"], source_key="image"),
>                 # randomly crop out patch samples from big image based on pos / neg ratio
>                 # the image centers of negative samples must be in valid image area
>                 RandCropByPosNegLabeld(
>                     keys=["image", "label"],
>                     label_key="label",
>                     spatial_size=(96, 96, 96),
>                     pos=1,
>                     neg=1,
>                     num_samples=4,
>                     image_key="image",
>                     image_threshold=0,
>                 ),
>                 # user can also add other random transforms
>                 # RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=1.0, spatial_size=(96, 96, 96),
>                 #             rotate_range=(0, 0, np.pi/15), scale_range=(0.1, 0.1, 0.1)),
>                 ToTensord(keys=["image", "label"]),
>             ]
>         )
>         val_transforms = Compose(
>             [
>                 LoadNiftid(keys=["image", "label"]),
>                 AddChanneld(keys=["image", "label"]),
>                 Spacingd(
>                     keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest"),
>                 ),
>                 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,
>                 ),
>                 #CropForegroundd(keys=["image", "label"], source_key="image"),
>                 ToTensord(keys=["image", "label"]),
>             ]
>         )
> 
>         # we use cached datasets - these are 10x faster than regular datasets
>         self.train_ds = CacheDataset(
>             data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4,
>         )
>         self.val_ds = CacheDataset(
>             data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4,
>         )
>         # self.train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
>         # self.val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
> 
>     def train_dataloader(self):
>         train_loader = torch.utils.data.DataLoader(
>             self.train_ds, batch_size=2, shuffle=True, num_workers=4, collate_fn=list_data_collate,
>         )
>         return train_loader
> 
>     def val_dataloader(self):
>         val_loader = torch.utils.data.DataLoader(self.val_ds, batch_size=1, num_workers=4)
>         return val_loader
> 
>     def configure_optimizers(self):
>         optimizer = torch.optim.Adam(self._model.parameters(), 1e-4)
>         return optimizer
> 
>     def training_step(self, batch, batch_idx):
>         images, labels = batch["image"], batch["label"]
>         output = self.forward(images)
>         loss = self.loss_function(output, labels)
>         tensorboard_logs = {"train_loss": loss.item()}
>         return {"loss": loss, "log": tensorboard_logs}
> 
>     def validation_step(self, batch, batch_idx):
>         images, labels = batch["image"], batch["label"]
>         roi_size = (160, 160, 160)
>         sw_batch_size = 4
>         outputs = sliding_window_inference(images, roi_size, sw_batch_size, self.forward)
>         loss = self.loss_function(outputs, labels)
>         value = compute_meandice(
>             y_pred=outputs,
>             y=labels,
>             include_background=False,
>             to_onehot_y=True,
>             mutually_exclusive=True,
>         )
>         return {"val_loss": loss, "val_dice": value}
> 
>     def validation_epoch_end(self, outputs):
>         val_dice, val_loss, num_items = 0, 0, 0
>         for output in outputs:
>             val_dice += output["val_dice"].sum().item()
>             val_loss += output["val_loss"].sum().item()
>             num_items += len(output["val_dice"])
>         mean_val_dice = torch.tensor(val_dice / num_items)
>         mean_val_loss = torch.tensor(val_loss / num_items)
>         tensorboard_logs = {
>             "val_dice": mean_val_dice,
>             "val_loss": mean_val_loss,
>         }
>         if mean_val_dice > self.best_val_dice:
>             self.best_val_dice = mean_val_dice
>             self.best_val_epoch = self.current_epoch
>         print(
>             f"current epoch: {self.current_epoch} current mean dice: {mean_val_dice:.4f}"
>             f"\nbest mean dice: {self.best_val_dice:.4f} at epoch: {self.best_val_epoch}"
>         )
>         return {"log": tensorboard_logs}
> net = Net()
> log_dir = os.path.join(root_dir, "logs")
> tb_logger = pytorch_lightning.loggers.TensorBoardLogger(
>     save_dir=log_dir
> )
> checkpoint_callback = pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint(
>     filepath=os.path.join(log_dir, "{epoch}-{val_loss:.2f}-{val_dice:.2f}")
> )
> 
> trainer = pytorch_lightning.Trainer(
>     gpus=[0,1],
>     max_epochs=600,
>     logger=tb_logger,
>     checkpoint_callback=checkpoint_callback,
>     num_sanity_val_steps=1,
> )
> 
> trainer.fit(net)
> 
> print(f"train completed, best_metric: {net.best_val_dice:.4f} at epoch {net.best_val_epoch}")
Lightning Segmentation tutorial breaks on multiple GPU's (CacheDataset?)
Reproducing the spleen Lightning segmentation tutorial from [here](https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/spleen_segmentation_3d_lightning.ipynb). Full logs of two errors (1) and (2) shared below. You can reproduce with full code from the notebook into a singly python file can be copied from (3)
With the latest library versions, and running on two GPU's results in the following error (1):
In multiprocessing/reduction.py > def dump:
![image](https://user-images.githubusercontent.com/10415569/96673444-95c9d300-1377-11eb-83c9-ddc161190f27.png)
After reading full log (1) I've commented out `CropForegroundd(keys=["image", "label"], source_key="image"),` from both `train_transforms` and `val_transforms` and I get the following error (2):
In multiprocessing/reduction.py > def dump:
![image](https://user-images.githubusercontent.com/10415569/96673846-8303ce00-1378-11eb-97e8-1956d7208782.png)



(1) Full log:

> Exception has occurred: PicklingError
> Can't pickle <function CropForegroundd.<lambda> at 0x7fd6e5159b80>: attribute lookup CropForegroundd.<lambda> on monai.transforms.croppad.dictionary failed
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
>     ForkingPickler(file, protocol).dump(obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
>     reduction.dump(process_obj, fp)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
>     self._launch(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
>     super().__init__(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
>     return Popen(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/process.py", line 121, in start
>     self._popen = self._Popen(self)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 149, in start_processes
>     process.start()
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
>     return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_spawn_backend.py", line 64, in train
>     mp.spawn(self.ddp_train, nprocs=self.nprocs, args=(self.mp_queue, model,))
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
>     results = self.accelerator_backend.train()
>   File "/nfs/users/ext_bart.michiels/projects/examples/test_multi_gpu.py", line 213, in <module>
>     trainer.fit(net)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 97, in _run_module_code
>     _run_code(code, mod_globals, init_globals,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 265, in run_path
>     return _run_module_code(code, init_globals, run_name,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 194, in _run_module_as_main
>     return _run_code(code, main_globals, None,

(2) Full log:

> Exception has occurred: TypeError
> cannot pickle '_thread.lock' object
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
>     ForkingPickler(file, protocol).dump(obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
>     reduction.dump(process_obj, fp)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
>     self._launch(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
>     super().__init__(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
>     return Popen(process_obj)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/multiprocessing/process.py", line 121, in start
>     self._popen = self._Popen(self)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 149, in start_processes
>     process.start()
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
>     return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_spawn_backend.py", line 64, in train
>     mp.spawn(self.ddp_train, nprocs=self.nprocs, args=(self.mp_queue, model,))
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 454, in fit
>     results = self.accelerator_backend.train()
>   File "/nfs/users/ext_bart.michiels/projects/examples/test_multi_gpu.py", line 213, in <module>
>     trainer.fit(net)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 97, in _run_module_code
>     _run_code(code, mod_globals, init_globals,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 265, in run_path
>     return _run_module_code(code, init_globals, run_name,
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 87, in _run_code
>     exec(code, run_globals)
>   File "/nfs/projects/dgx1-pod1-medical-project/anaconda/envs/covid_clf/lib/python3.8/runpy.py", line 194, in _run_module_as_main
>     return _run_code(code, main_globals, None,

(3) Full python code from the notebook to reproduce the errors:

> import glob
> import os
> import shutil
> import tempfile
> import matplotlib.pyplot as plt
> import pytorch_lightning
> import torch
> from monai.apps import download_and_extract
> from monai.config import print_config
> from monai.data import CacheDataset, list_data_collate
> from monai.inferers import sliding_window_inference
> from monai.losses import DiceLoss
> from monai.metrics import compute_meandice
> from monai.networks.layers import Norm
> from monai.networks.nets import UNet
> from monai.transforms import (
>     AddChanneld,
>     Compose,
>     CropForegroundd,
>     LoadNiftid,
>     Orientationd,
>     RandCropByPosNegLabeld,
>     ScaleIntensityRanged,
>     Spacingd,
>     ToTensord,
> )
> from monai.utils import set_determinism
> print_config()
> directory = os.environ.get("MONAI_DATA_DIRECTORY")
> root_dir = tempfile.mkdtemp() if directory is None else directory
> print(root_dir)
> resource = "https://drive.google.com/uc?id=1jzeNU1EKnK81PyTsrx0ujfNl-t0Jo8uE"
> md5 = "410d4a301da4e5b2f6f86ec3ddba524e"
> compressed_file = os.path.join(root_dir, "Task09_Spleen.tar")
> data_dir = os.path.join(root_dir, "Task09_Spleen")
> if not os.path.exists(data_dir):
>     download_and_extract(resource, compressed_file, root_dir, md5)
> class Net(pytorch_lightning.LightningModule):
>     def __init__(self):
>         super().__init__()
>         self._model = UNet(
>             dimensions=3,
>             in_channels=1,
>             out_channels=2,
>             channels=(16, 32, 64, 128, 256),
>             strides=(2, 2, 2, 2),
>             num_res_units=2,
>             norm=Norm.BATCH,
>         )
>         self.loss_function = DiceLoss(to_onehot_y=True, softmax=True)
>         self.best_val_dice = 0
>         self.best_val_epoch = 0
> 
>     def forward(self, x):
>         return self._model(x)
> 
>     def prepare_data(self):
>         # set up the correct data path
>         train_images = sorted(glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
>         train_labels = sorted(glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
>         data_dicts = [
>             {"image": image_name, "label": label_name}
>             for image_name, label_name in zip(train_images, train_labels)
>         ]
>         train_files, val_files = data_dicts[:-9], data_dicts[-9:]
> 
>         # set deterministic training for reproducibility
>         set_determinism(seed=0)
> 
>         # define the data transforms
>         train_transforms = Compose(
>             [
>                 LoadNiftid(keys=["image", "label"]),
>                 AddChanneld(keys=["image", "label"]),
>                 Spacingd(
>                     keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest"),
>                 ),
>                 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,
>                 ),
>                 #CropForegroundd(keys=["image", "label"], source_key="image"),
>                 # randomly crop out patch samples from big image based on pos / neg ratio
>                 # the image centers of negative samples must be in valid image area
>                 RandCropByPosNegLabeld(
>                     keys=["image", "label"],
>                     label_key="label",
>                     spatial_size=(96, 96, 96),
>                     pos=1,
>                     neg=1,
>                     num_samples=4,
>                     image_key="image",
>                     image_threshold=0,
>                 ),
>                 # user can also add other random transforms
>                 # RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=1.0, spatial_size=(96, 96, 96),
>                 #             rotate_range=(0, 0, np.pi/15), scale_range=(0.1, 0.1, 0.1)),
>                 ToTensord(keys=["image", "label"]),
>             ]
>         )
>         val_transforms = Compose(
>             [
>                 LoadNiftid(keys=["image", "label"]),
>                 AddChanneld(keys=["image", "label"]),
>                 Spacingd(
>                     keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest"),
>                 ),
>                 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,
>                 ),
>                 #CropForegroundd(keys=["image", "label"], source_key="image"),
>                 ToTensord(keys=["image", "label"]),
>             ]
>         )
> 
>         # we use cached datasets - these are 10x faster than regular datasets
>         self.train_ds = CacheDataset(
>             data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4,
>         )
>         self.val_ds = CacheDataset(
>             data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4,
>         )
>         # self.train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
>         # self.val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
> 
>     def train_dataloader(self):
>         train_loader = torch.utils.data.DataLoader(
>             self.train_ds, batch_size=2, shuffle=True, num_workers=4, collate_fn=list_data_collate,
>         )
>         return train_loader
> 
>     def val_dataloader(self):
>         val_loader = torch.utils.data.DataLoader(self.val_ds, batch_size=1, num_workers=4)
>         return val_loader
> 
>     def configure_optimizers(self):
>         optimizer = torch.optim.Adam(self._model.parameters(), 1e-4)
>         return optimizer
> 
>     def training_step(self, batch, batch_idx):
>         images, labels = batch["image"], batch["label"]
>         output = self.forward(images)
>         loss = self.loss_function(output, labels)
>         tensorboard_logs = {"train_loss": loss.item()}
>         return {"loss": loss, "log": tensorboard_logs}
> 
>     def validation_step(self, batch, batch_idx):
>         images, labels = batch["image"], batch["label"]
>         roi_size = (160, 160, 160)
>         sw_batch_size = 4
>         outputs = sliding_window_inference(images, roi_size, sw_batch_size, self.forward)
>         loss = self.loss_function(outputs, labels)
>         value = compute_meandice(
>             y_pred=outputs,
>             y=labels,
>             include_background=False,
>             to_onehot_y=True,
>             mutually_exclusive=True,
>         )
>         return {"val_loss": loss, "val_dice": value}
> 
>     def validation_epoch_end(self, outputs):
>         val_dice, val_loss, num_items = 0, 0, 0
>         for output in outputs:
>             val_dice += output["val_dice"].sum().item()
>             val_loss += output["val_loss"].sum().item()
>             num_items += len(output["val_dice"])
>         mean_val_dice = torch.tensor(val_dice / num_items)
>         mean_val_loss = torch.tensor(val_loss / num_items)
>         tensorboard_logs = {
>             "val_dice": mean_val_dice,
>             "val_loss": mean_val_loss,
>         }
>         if mean_val_dice > self.best_val_dice:
>             self.best_val_dice = mean_val_dice
>             self.best_val_epoch = self.current_epoch
>         print(
>             f"current epoch: {self.current_epoch} current mean dice: {mean_val_dice:.4f}"
>             f"\nbest mean dice: {self.best_val_dice:.4f} at epoch: {self.best_val_epoch}"
>         )
>         return {"log": tensorboard_logs}
> net = Net()
> log_dir = os.path.join(root_dir, "logs")
> tb_logger = pytorch_lightning.loggers.TensorBoardLogger(
>     save_dir=log_dir
> )
> checkpoint_callback = pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint(
>     filepath=os.path.join(log_dir, "{epoch}-{val_loss:.2f}-{val_dice:.2f}")
> )
> 
> trainer = pytorch_lightning.Trainer(
>     gpus=[0,1],
>     max_epochs=600,
>     logger=tb_logger,
>     checkpoint_callback=checkpoint_callback,
>     num_sanity_val_steps=1,
> )
> 
> trainer.fit(net)
> 
> print(f"train completed, best_metric: {net.best_val_dice:.4f} at epoch {net.best_val_epoch}")


Discussion:
Hi @bartmch ,

Thanks for your cool experiments and feedback.
If you train with distributed data parallel, I think you need to use `DistributedSampler` for the dataset? Check this example:
https://github.com/Project-MONAI/tutorials/blob/master/acceleration/distributed_training/unet_training_ddp.py#L117
I am not very familiar with PyTorch Lightning, maybe @marksgraham can help explain more details.

Thanks.

Hi @bartmch ,

Thanks for your cool experiments and feedback.
If you train with distributed data parallel, I think you need to use `DistributedSampler` for the dataset? Check this example:
https://github.com/Project-MONAI/tutorials/blob/master/acceleration/distributed_training/unet_training_ddp.py#L117
I am not very familiar with PyTorch Lightning, maybe @marksgraham can help explain more details.

Thanks.


Resolving patch:
diff --git a/monai/data/dataset.py b/monai/data/dataset.py
--- a/monai/data/dataset.py
+++ b/monai/data/dataset.py
@@ -219,8 +219,7 @@ def _pre_first_random_cachecheck(self, item_transformed):
 
     def __getitem__(self, index: int):
         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
+        return self._first_random_and_beyond_transform(pre_random_item)
 
 
 class LMDBDataset(PersistentDataset):
@@ -276,16 +275,10 @@ def __init__(
         self.lmdb_kwargs = lmdb_kwargs or {}
         if not self.lmdb_kwargs.get("map_size", 0):
             self.lmdb_kwargs["map_size"] = 1024 ** 4  # default map_size
+        self._read_env = None
 
+    def _fill_cache_start_reader(self):
         # create cache
-        self._fill_cache()
-        # read-only database env
-        self.lmdb_kwargs["readonly"] = True
-        if self.lmdb_kwargs.get("lock", None) is None:
-            self.lmdb_kwargs["lock"] = False
-        self._read_env = lmdb.open(path=f"{self.db_file}", subdir=False, **self.lmdb_kwargs)
-
-    def _fill_cache(self):
         print(f"Accessing lmdb file: {self.db_file.absolute()}.")
         self.lmdb_kwargs["readonly"] = False
         env = lmdb.open(path=f"{self.db_file}", subdir=False, **self.lmdb_kwargs)
@@ -316,12 +309,23 @@ def _fill_cache(self):
                 size = env.info()["map_size"]
                 env.close()
                 raise ValueError(f"LMDB map size reached, increase size above current size of {size}.")
+        size = env.info()["map_size"]
         env.close()
+        # read-only database env
+        self.lmdb_kwargs["readonly"] = True
+        self.lmdb_kwargs["map_size"] = size
+        if self.lmdb_kwargs.get("lock", None) is None:
+            self.lmdb_kwargs["lock"] = False
+        if self.lmdb_kwargs.get("readahead", None) is None:
+            self.lmdb_kwargs["readahead"] = False
+        return lmdb.open(path=f"{self.db_file}", subdir=False, **self.lmdb_kwargs)
 
     def _pre_first_random_cachecheck(self, item_transformed):
         """
         if the item is not found in the lmdb file, resolves to the persistent cache default behaviour.
         """
+        if self._read_env is None:
+            self._read_env = self._fill_cache_start_reader()
         with self._read_env.begin(write=False) as txn:
             data = txn.get(self.hash_func(item_transformed))
         if data is None:
@@ -336,6 +340,8 @@ def info(self):
         """
         Returns: dataset info dictionary.
         """
+        if self._read_env is None:
+            self._read_env = self._fill_cache_start_reader()
         out = dict(self._read_env.info())
         out["size"] = len(self.data)
         out["filename"] = f"{self.db_file.absolute()}"
@@ -382,7 +388,7 @@ def __init__(
         transform: Union[Sequence[Callable], Callable],
         cache_num: int = sys.maxsize,
         cache_rate: float = 1.0,
-        num_workers: int = 0,
+        num_workers: Optional[int] = None,
     ) -> None:
         """
         Args:
@@ -392,79 +398,52 @@ 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.
-                If 0 a single thread will be used. Default is 0.
-        """
-        if not isinstance(transform, Compose):
-            transform = Compose(transform)
-        super().__init__(data, transform)
-        self.cache_num = min(cache_num, int(len(data) * cache_rate), len(data))
-        if self.cache_num > 0:
-            self._cache = [None] * self.cache_num
+            num_workers: the number of worker processes to use.
+                If num_workers is None then the number returned by os.cpu_count() is used.
+        """
+        super().__init__(data=data, transform=transform if isinstance(transform, Compose) else Compose(transform))
+        self.cache_num = min(int(cache_num), int(len(data) * cache_rate), len(data))
+        self.num_workers = num_workers
+        if self.num_workers is not None:
+            self.num_workers = max(int(self.num_workers), 1)
+        self._cache: List = self._fill_cache()
+
+    def _fill_cache(self) -> List:
+        if self.cache_num <= 0:
+            return []
+        if not has_tqdm:
+            warnings.warn("tqdm is not installed, will not show the caching progress bar.")
+        with ThreadPool(self.num_workers) as p:
             if has_tqdm:
-                pbar = tqdm(total=self.cache_num, desc="Load and cache transformed data")
-            else:
-                warnings.warn("tqdm is not installed, will not show the caching progress bar.")
-                pbar = None
-
-            if num_workers > 0:
-                self._item_processed = 0
-                self._thread_lock = threading.Lock()
-                with ThreadPool(num_workers) as p:
-                    p.map(
-                        self._load_cache_item_thread,
-                        [(i, data[i], transform.transforms, pbar) for i in range(self.cache_num)],
-                    )
-            else:
-                for i in range(self.cache_num):
-                    self._cache[i] = self._load_cache_item(data[i], transform.transforms)
-                    if pbar is not None:
-                        pbar.update(1)
-            if pbar is not None:
-                pbar.close()
+                return list(tqdm(p.imap(self._load_cache_item, range(self.cache_num)), total=self.cache_num))
+            return list(p.imap(self._load_cache_item, range(self.cache_num)))
 
-    def _load_cache_item(self, item: Any, transforms: Sequence[Callable]):
+    def _load_cache_item(self, idx: int):
         """
         Args:
-            item: input item to load and transform to generate dataset for model.
-            transforms: transforms to execute operations on input item.
+            idx: the index of the input data sequence.
         """
-        for _transform in transforms:
+        item = self.data[idx]
+        for _transform in self.transform.transforms:  # type: ignore
             # execute all the deterministic transforms
             if isinstance(_transform, Randomizable) or not isinstance(_transform, Transform):
                 break
             item = apply_transform(_transform, item)
         return item
 
-    def _load_cache_item_thread(self, args: Any) -> None:
-        """
-        Args:
-            args: tuple with contents (i, item, transforms, pbar).
-                i: the index to load the cached item to.
-                item: input item to load and transform to generate dataset for model.
-                transforms: transforms to execute operations on input item.
-                pbar: tqdm progress bar
-        """
-        i, item, transforms, pbar = args
-        self._cache[i] = self._load_cache_item(item, transforms)
-        if pbar is not None:
-            with self._thread_lock:
-                pbar.update(1)
-
     def __getitem__(self, index):
-        if index < self.cache_num:
-            # load data from cache and execute from the first random transform
-            start_run = False
-            data = self._cache[index]
-            for _transform in self.transform.transforms:  # pytype: disable=attribute-error
-                if not start_run and not isinstance(_transform, Randomizable) and isinstance(_transform, Transform):
-                    continue
-                else:
-                    start_run = True
+        if index >= self.cache_num:
+            # no cache for this index, execute all the transforms directly
+            return super(CacheDataset, self).__getitem__(index)
+        # load data from cache and execute from the first random transform
+        start_run = False
+        if self._cache is None:
+            self._cache = self._fill_cache()
+        data = self._cache[index]
+        for _transform in self.transform.transforms:  # pytype: disable=attribute-error
+            if start_run or isinstance(_transform, Randomizable) or not isinstance(_transform, Transform):
+                start_run = True
                 data = apply_transform(_transform, data)
-        else:
-            # no cache for this data, execute all the transforms directly
-            data = super(CacheDataset, self).__getitem__(index)
         return data
 
 
@@ -508,7 +487,7 @@ def __init__(
         replace_rate: float,
         cache_num: int = sys.maxsize,
         cache_rate: float = 1.0,
-        num_init_workers: int = 0,
+        num_init_workers: Optional[int] = None,
         num_replace_workers: int = 0,
     ) -> None:
         """
@@ -521,11 +500,13 @@ def __init__(
             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.
+                If num_init_workers is None then the number returned by os.cpu_count() is used.
             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.
         """
         super().__init__(data, transform, cache_num, cache_rate, num_init_workers)
+        if self._cache is None:
+            self._cache = self._fill_cache()
         if self.cache_num >= len(data):
             raise ValueError("cache_num must be smaller than dataset length to support replacement.")
         if replace_rate <= 0:
@@ -588,25 +569,25 @@ def _try_update_cache(self):
 
         """
         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]
+            if not self._replace_done:
+                return False
 
-                self._start_pos += self._replace_num
-                if self._start_pos >= self._total_num:
-                    self._start_pos -= self._total_num
+            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._compute_data_idx()
+            self._start_pos += self._replace_num
+            if self._start_pos >= self._total_num:
+                self._start_pos -= self._total_num
 
-                # ready for next round
-                self._round += 1
-                self._replace_done = False
-                return True
-            else:
-                return False
+            self._compute_data_idx()
+
+            # ready for next round
+            self._round += 1
+            self._replace_done = False
+            return True
 
     def update_cache(self):
         """
@@ -653,7 +634,7 @@ def _replace_cache_thread(self, index: int):
 
         """
         pos: int = self._replace_data_idx[index]
-        self._replacements[index] = self._load_cache_item(self.data[pos], self.transform.transforms)  # type: ignore
+        self._replacements[index] = self._load_cache_item(pos)
 
     def _compute_replacements(self):
         """


Test patch: