ddp, cache on GPU and GPU transforms

[hw-ju]

Hi!
I'm adapting this tutorial script to run in ddp https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/unetr_btcv_segmentation_3d.ipynb.

It works when I use CacheDataset for CPU mem caching (together with DistributedSampler as arguments to DataLoader). However, after adding EnsureTyped and ToDeviced to cache on GPU and perform transforms on GPU (I checked that those transforms support tensor as input by python monai/transforms/utils.py) as

 train_transforms = Compose(
        [
            LoadImaged(keys=["image", "label"]),
            AddChanneld(keys=["image", "label"]),
            Orientationd(keys=["image", "label"], axcodes="RAS"),
            Spacingd(
                keys=["image", "label"],
                pixdim=(1.5, 1.5, 2.0),
                mode=("bilinear", "nearest"),
            ),
            ScaleIntensityRanged(
                keys=["image"],
                a_min=-175,
                a_max=250,
                b_min=0.0,
                b_max=1.0,
                clip=True,
            ),
            CropForegroundd(keys=["image", "label"], source_key="image"),
            EnsureTyped(keys=["image", "label"], data_type="tensor"), 
            ToDeviced(keys=["image", "label"], device=f"cuda:{args.local_rank}"),            
            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,
            ),
            RandFlipd(
                keys=["image", "label"],
                spatial_axis=[0],
                prob=0.10,
            ),
            RandFlipd(
                keys=["image", "label"],
                spatial_axis=[1],
                prob=0.10,
            ),
            RandFlipd(
                keys=["image", "label"],
                spatial_axis=[2],
                prob=0.10,
            ),
            RandRotate90d(
                keys=["image", "label"],
                prob=0.10,
                max_k=3,
            ),
            RandShiftIntensityd(
                keys=["image"],
                offsets=0.10,
                prob=0.50,
            ),
            ToTensord(keys=["image", "label"]),
        ]
    )

    val_transforms = Compose(
        [
            LoadImaged(keys=["image", "label"]),
            AddChanneld(keys=["image", "label"]),
            Orientationd(keys=["image", "label"], axcodes="RAS"),
            Spacingd(
                keys=["image", "label"],
                pixdim=(1.5, 1.5, 2.0),
                mode=("bilinear", "nearest"),
            ),
            ScaleIntensityRanged(
                keys=["image"], a_min=-175, a_max=250, b_min=0.0, b_max=1.0, clip=True
            ),
            CropForegroundd(keys=["image", "label"], source_key="image"),
            ToTensord(keys=["image", "label"]),
            EnsureTyped(keys=["image", "label"], data_type="tensor"), 
            ToDeviced(keys=["image", "label"], device=f"cuda:{args.local_rank}"),            
        ]
    )

I get error:

Traceback (most recent call last):
  File "/home/u00u7uv9zk2dwK914t357/run_monaicore/multigpu/unetr_ddp/gputransform/unetr_btcv_ddp_gputransform.py", line 308, in <module>
    main()
  File "/home/u00u7uv9zk2dwK914t357/run_monaicore/multigpu/unetr_ddp/gputransform/unetr_btcv_ddp_gputransform.py", line 304, in main
    train(args=args)
  File "/home/u00u7uv9zk2dwK914t357/run_monaicore/multigpu/unetr_ddp/gputransform/unetr_btcv_ddp_gputransform.py", line 225, in train
    for batch_data in train_loader:
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1203, in _next_data
    return self._process_data(data)
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1229, in _process_data
    data.reraise()
  File "/opt/conda/lib/python3.8/site-packages/torch/_utils.py", line 434, in reraise
    raise exception
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/opt/monai/monai/data/dataset.py", line 97, in __getitem__
    return self._transform(index)
  File "/opt/monai/monai/data/dataset.py", line 805, in _transform
    data = deepcopy(data)
  File "/opt/conda/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/opt/conda/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/opt/conda/lib/python3.8/copy.py", line 153, in deepcopy
    y = copier(memo)
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 97, in __deepcopy__
    new_storage = self.storage().__deepcopy__(memo)
  File "/opt/conda/lib/python3.8/site-packages/torch/storage.py", line 48, in __deepcopy__
    new_storage = self.clone()
  File "/opt/conda/lib/python3.8/site-packages/torch/storage.py", line 63, in clone
    with torch.cuda.device(device):
  File "/opt/conda/lib/python3.8/site-packages/torch/cuda/__init__.py", line 272, in __enter__
    self.prev_idx = torch.cuda.current_device()
  File "/opt/conda/lib/python3.8/site-packages/torch/cuda/__init__.py", line 479, in current_device
    _lazy_init()
  File "/opt/conda/lib/python3.8/site-packages/torch/cuda/__init__.py", line 204, in _lazy_init
    raise RuntimeError(
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

Any thoughts?

[Nic-Ma]

Hi @hw-ju ,

CUDA program can't support forked multi-processing, so you should not use num_workers > 0 of DataLoader.
Suggest you to try ThreadDataLoader instead:
https://github.com/Project-MONAI/tutorials/blob/master/acceleration/distributed_training/brats_training_ddp.py#L220
For more details please refer to:
https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_model_training_guide.md

Thanks.

[hw-ju]

Thanks @Nic-Ma!
So for GPU transforms, we can either use ThreadDataLoader or DataLoader with num_worker=0?

[Nic-Ma]

Yes, ThreadDataLoader may have better performance.

Thanks.