Best practice for using RandTransforms to reduce time complexity

[Zrrr1997]

Dear MONAI team,

first of all, thanks for your wonderful work and support! I am currently working on a PET/CT dataset and am observing a dramatic change in the inference/training time when I add Random transforms to my pipeline.

Details about my configuration
The PET and CT volumes are registered and have an average size of (400, 400, 128) and there are around 400 training volumes. I am using an A100 NVIDIA GPU.

My training transforms pipeline is as follows:

            train_transforms = Compose(
                [
                    LoadImaged(keys=["ct_vol", "pet_vol", "seg"]),
                    AddChanneld(keys=["ct_vol", "pet_vol", "seg"]),
                    Spacingd(keys=["ct_vol", "pet_vol", "seg"], pixdim=pixdim, mode=("bilinear", "bilinear", "nearest")),
                    Orientationd(keys=["ct_vol", "pet_vol", "seg"], axcodes="LAS"),
                    ScaleIntensityRanged(keys=["ct_vol"], a_min=a_min_ct, a_max=a_max_ct, b_min=0.0, b_max=1.0, clip=True),
                    ScaleIntensityRanged(keys=["pet_vol"], a_min=a_min_pet, a_max=a_max_pet, b_min=0.0, b_max=1.0, clip=True),
                    CropForegroundd(keys=["ct_vol", "pet_vol", "seg"], source_key="pet_vol"),
                    RandRotated(keys=["ct_vol", "pet_vol", "seg"], range_x=np.pi / 6, range_y=np.pi / 6, range_z=np.pi / 6, prob=0.2),
                    RandAffined(keys=["ct_vol", "pet_vol", "seg"], prob=0.2, scale_range=[-0.3, 0.4]),
                    RandFlipd(keys=["ct_vol", "pet_vol", "seg"], prob=0.5, spatial_axis=0),
                    RandFlipd(keys=["ct_vol", "pet_vol", "seg"], prob=0.5, spatial_axis=1),
                    RandFlipd(keys=["ct_vol", "pet_vol", "seg"], prob=0.5, spatial_axis=2),
                    Resized(keys=["ct_vol", "pet_vol"], spatial_size=spatial_size),
                    Resized(keys=["seg"], spatial_size=spatial_size, mode="nearest-exact"),
                    ConcatItemsd(keys=["ct_vol", "pet_vol"], name="ct_pet_vol", dim=0),  # concatenate pet and ct channels
                    ToTensord(keys=["ct_vol", "pet_vol", "seg", "ct_pet_vol", "class_label"]),
                ]
            )

Using these transforms leads to a large time overhead - around 45-75 minutes for a single epoch. However, removing all RandFlip, RandAffine, RandRotate transforms leads to epoch times of only 6 minutes. Additionally, I am using a CacheDataset to load the data and a standard Inferer (no sliding window applied).

I was wondering what are the best practices when using random transforms. Am I using the wrong Dataset class? Is the order of the transforms not optimal? Or are such inference/training times normal for such heavy transforms + amount of data.

I would gladly read more about this, or would be happy if anyone can point me in the right direction!

Best regards,
Zdravko

[KumoLiu]

Hi @Zrrr1997,
Since CacheDataset will cache the outcomes before the first Randomizable Transform within a Compose instance. So to improve the caching efficiency, please always put as many as possible non-random transforms before the randomized ones when composing the chain of transforms. In your case, maybe you can move Resized before Randomizable Transform.
You can also use AMP(auto mixed precision) during your training which can significantly improved training speed.
For more tips you can refer to fast_tranining_tutorial.
Hope it can helps you!

[Zrrr1997]

Hi @KumoLiu,
thanks for the answer! It seems that moving the Resized transforms and adding a EnsureType(...., device=torch.device('cuda:0') before all Randomizable Transforms and using ThreadDataLoader leads to a substantial improvement in the time complexity.

However, using CacheDataset is then quite difficult to do, since the dataset is quite large when you pile up the randomizable transforms and cannot be cached. I also set the copy_cache=False, which led to some improvement, but I cannot cache the whole dataset yet.

What would you propose in this scenario? I tried PersistentDataset and it works fine, but it needs quite a lot of time for the first epoch.

P.S. I have 4 A100 GPUs at my disposal. Is it possible to somehow distribute the cache in a multi-GPU setup with the CacheDataset?

[KumoLiu]

For multigpu, you can get some instructions refer to this tutorial https://github.com/Project-MONAI/tutorials/blob/main/acceleration/multi_gpu_test.ipynb.
@wyli @dongyang0122, can you share more example or use case? Thanks in advance!