Combining SmartCacheDataset with Pytorch Lightning multi-gpu training

[jxchen01]

I am trying to combine SmartCacheDataset with Pytorch Lightning so that I can take advantage of all the training tricks, like mutli-gpu trainings, or stochastic weight averaging, etc..

I adapted the 3d_segmentation_lightning tutorial online, into the following script:

import pytorch_lightning
from monai.utils import set_determinism
from monai.transforms import (
    AsDiscrete,
    AddChanneld,
    Compose,
    CropForegroundd,
    LoadImaged,
    Orientationd,
    RandCropByPosNegLabeld,
    ScaleIntensityRanged,
    Spacingd,
    EnsureTyped,
    EnsureType,
)
from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.metrics import DiceMetric
from monai.losses import DiceLoss
from monai.inferers import sliding_window_inference
from monai.data import SmartCacheDataset, CacheDataset, Dataset, list_data_collate, decollate_batch
from monai.config import print_config
from monai.apps import download_and_extract
import torch
import matplotlib.pyplot as plt
import tempfile
import shutil
import os
import glob


class myDataModule(pytorch_lightning.LightningDataModule):
    def __init__(self):
        super().__init__()

    def prepare_data(self):
        root_dir = "./tmp/"
        data_dir = os.path.join(root_dir, "Task09_Spleen")
        # 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)
        ]
        print("total data: ..............")
        print(len(data_dicts))
        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(
            [
                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=-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)),
                EnsureTyped(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=-57, a_max=164,
                    b_min=0.0, b_max=1.0, clip=True,
                ),
                CropForegroundd(keys=["image", "label"], source_key="image"),
                EnsureTyped(keys=["image", "label"]),
            ]
        )

        # we use cached datasets - these are 10x faster than regular datasets
        self.train_ds = SmartCacheDataset(
            data=train_files, transform=train_transforms,
            replace_rate=0.2, num_replace_workers=4,
            cache_rate=0.5, num_init_workers=4
        )
        self.val_ds = SmartCacheDataset(
            data=val_files, transform=val_transforms,
            replace_rate=0.2, num_replace_workers=4,
            cache_rate=0.5, num_init_workers=4,
        )
        """
        self.train_ds = Dataset(
            data=train_files, transform=train_transforms)
        self.val_ds = 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=2, collate_fn=list_data_collate, persistent_workers=False
        )
        return train_loader

    def val_dataloader(self):
        val_loader = torch.utils.data.DataLoader(
            self.val_ds, batch_size=1, num_workers=2, persistent_workers=False)
        return val_loader


class Net(pytorch_lightning.LightningModule):
    def __init__(self):
        super().__init__()
        self._model = UNet(
            spatial_dims=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.post_pred = Compose([EnsureType("tensor", device="cpu"), AsDiscrete(argmax=True, to_onehot=2)])
        self.post_label = Compose([EnsureType("tensor", device="cpu"), AsDiscrete(to_onehot=2)])
        self.dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)
        self.best_val_dice = 0
        self.best_val_epoch = 0

    def forward(self, x):
        return self._model(x)

    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)
        outputs = [self.post_pred(i) for i in decollate_batch(outputs)]
        labels = [self.post_label(i) for i in decollate_batch(labels)]
        self.dice_metric(y_pred=outputs, y=labels)
        return {"val_loss": loss, "val_number": len(outputs)}

    def validation_epoch_end(self, outputs):
        val_loss, num_items = 0, 0
        for output in outputs:
            val_loss += output["val_loss"].sum().item()
            num_items += output["val_number"]
        mean_val_dice = self.dice_metric.aggregate().item()
        self.dice_metric.reset()
        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} "
            f"current mean dice: {mean_val_dice:.4f}"
            f"\nbest mean dice: {self.best_val_dice:.4f} "
            f"at epoch: {self.best_val_epoch}"
        )
        return {"log": tensorboard_logs}


if __name__ == '__main__':
    print(torch.multiprocessing.get_start_method())
    # initialise the LightningModule
    net = Net()
    dm = myDataModule()
    dm.prepare_data()

    # set up loggers and checkpoints
    root_dir = "./tmp/"
    log_dir = os.path.join(root_dir, "logs")
    tb_logger = pytorch_lightning.loggers.TensorBoardLogger(
        save_dir=log_dir
    )

    # initialise Lightning's trainer.
    trainer = pytorch_lightning.Trainer(
        gpus=4,
        precision=16,
        max_epochs=600,
        logger=tb_logger,
        enable_checkpointing=True,
        num_sanity_val_steps=1,
        log_every_n_steps=16,
    )

    # train
    trainer.fit(net, datamodule=dm)

But, I kept getting this errors:

Traceback (most recent call last):
  File "test.py", line 240, in <module>
    trainer.fit(net, datamodule=dm)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 768, in fit
    self._call_and_handle_interrupt(
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 719, in _call_and_handle_interrupt
    return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/spawn.py", line 78, in launch
    mp.spawn(
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 179, in start_processes
    process.start()
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/im2im/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
TypeError: cannot pickle '_thread.lock' object

If I understand correctly, the error indicates that SmartCacheDataset is trying to dump some unpickable objects. Is that correct?

Thanks for the help in advance!
Jianxu

[Nic-Ma]

Hi @jxchen01 ,

Thanks for your interest here, could you please help confirm whether some of below items match your condition:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/dataset.py#L848

Thanks.

[jxchen01]

hmm... interesting... torchrun indeed works. I was not even aware of pytorch lightning needs to run with torchrun.

https://github.com/MMV-Lab/demo_box/blob/main/monai_w_plt_multigpu.py
I made a short script to reproduce this results.

python monai_w_plt_multigpu.py will not work. But, torchrun monai_w_plt_multigpu.py just works. This example only uses some random fake data. I will double check with some real data to verify if there is any difference when using torchrun.

Thanks so much for your help @Nic-Ma

[Nic-Ma]

Hi @jxchen01 ,

Thanks for your test feedback.
That makes sense, because your previous test is based on spawn method of multi-processing GPUs:
torch/multiprocessing/spawn.py

Thanks.

[Nic-Ma]

Thanks for your feedback, I enhanced the doc-string of SmartCacheDataset in PR: #4419.

Thanks.

[jxchen01]

@Nic-Ma , I did further investigation. It seems like torchrun makes the mutligpu script executable, but not in the expected way. After the code is running, only one GPU is being used. In other words, if I understand correctly, somehow torchrun dynamically scale the execution down to single GPU. I will do more exploration today. If you have any suggestion, I would love to hear more ideas.

[jxchen01]

it looks like even with the basic CacheDataset, torchrun still triggers only one GPU. I need to figure out how to use pytorch-lightning with torchrun properly.

[jxchen01]

@Nic-Ma After another day of exploration and learning, I think there might still be some issue with SmartCacheDataset (may not be a real issue, but some incompatibility as a plug-and-play module with other multi-node / multi-gpu training packages).

I did some more testing with Ray and Ray-lightning (https://github.com/ray-project/ray_lightning), and fairscale (https://github.com/facebookresearch/fairscale). One thing I learned is that the model object has to be serializable/pickleable. For example, here https://pytorch-lightning.readthedocs.io/en/1.4.0/advanced/multi_gpu.html#make-models-pickleable.

When I went into the code of SmartCacheDataset, I find this one: https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/dataset.py#L903. I think this might be the cause of the error message TypeError: cannot pickle '_thread.lock' object.

Just my 2 cents, not sure if my understanding is correct.

[Nic-Ma]

Hi @jxchen01 ,

The thread lock should not work with pickle in some cases.

Thanks.