Error with CacheDataset when Num_Workers > 1 (non-thread safe transforms)

[fogunsan]

Hi there, I'm wondering if I can get some assistance. I'm training a unet with paired images that have been randomly cropped to 32x32x32 from larger images that I stored on ram using cachedataset. I then use a PatchDataset implementing randcropbyposneglabeld and then CacheDataset to make training faster. However, the patch loading appears to be a major bottleneck with around 42000 patches (1000 patches for 42 images) taking around 15 hours to load. When I try to increase the number of workers from 1, the image pairs appear unaligned after plotting and I get this error in training:

RuntimeError: stack expects each tensor to be equal size, but got [1, 32, 32, 32] at entry 0 and [1, 32, 32, 25] at entry 2
Collate error on the key 'image' of dictionary data.

MONAI hint: if your transforms intentionally create images of different shapes, creating your DataLoader with collate_fn=pad_list_data_collate might solve this problem (check its documentation).

From the discussions I've read in the issues, I believe it is due to the randcropbyposneglabel transform and other random transforms I have added not yet being implemented as thread safe. I'm wondering if anyone has found a way around this problem to speed up caching rate. Is using PyTorch's thread safe random transforms and implementing them to be applied on both the image and label the only current solution?

Here is my code:

load_images= Compose(
    [
     	LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityd(keys = ["image"], minv=0.0, maxv=1.0)])
    
train_imgs_cache = CacheDataset(data=train_files, transform=load_images, cache_rate=1.0, num_workers = 1)
validate_imgs_cache = CacheDataset(data=validate_files, transform=load_images, cache_rate=1.0,num_workers =1)

patching_func= RandCropByPosNegLabeld( 
            keys = ["image", "label"],
            label_key = "image",
            spatial_size=(32,32,32),
            pos = 1,
            neg = 0.01, 
            num_samples= num_patches
        )
patch_transforms = Compose([
    RandRotated(keys =["image", "label"], range_x = 0.8, range_y = 0.8, range_z = 0.8, prob = 0.4), 
    RandFlipd(keys =["image", "label"], prob = 0.2, spatial_axis=1)])

train_patches_dataset = PatchDataset(data = train_imgs_cache, patch_func=patching_func, samples_per_image=num_patches, transform= patch_transforms)
validate_patches_dataset = PatchDataset(data = validate_imgs_cache, patch_func=patching_func, samples_per_image=num_patches, transform= patch_transfo>

train_patches_dataset = CacheDataset(data=train_patches_dataset,cache_rate=1.0, num_workers =1) 
validate_patches_dataset = CacheDataset(data=validate_patches_dataset, cache_rate=1.0, num_workers =1)

[KumoLiu]

Hi @fogunsan, from the error message, I think the error is due to your different sizes when stack them into a batch. Maybe you can try to add ResizeWithPadOrCropD in your patching_func to ensure that you have the same size.
Hope it can help you, thanks!

[fogunsan]

Hi @KumoLiu, thank you for your response. Upon adding ResizeWithPadOrCropD to my patching_func, it does silence the error message, however when plotting my image and corresponding target image in the training loop, they still do not appear to be aligned. I do not experience this issue when my num_workers is set to 1. Here's the modified code:

patching_func= Compose([RandCropByPosNegLabeld(
            keys = ["image", "label"],
            label_key = "image",
            spatial_size=(32,32,32),
            pos = 1,
            neg = 0.01, # 1
            num_samples= num_patches
        ), ResizeWithPadOrCropd(keys = ["image", "label"], spatial_size =(32,32,32))])

[KumoLiu]

Hi @fogunsan, I can't reproduce your problem, the code I used is shown below:

import numpy as np

import monai
import matplotlib.pyplot as plt
from monai.data import PatchDataset, DataLoader, CacheDataset
from monai.transforms import RandCropByPosNegLabeld, RandShiftIntensityd, Compose, LoadImaged, OrientationD
from monai.visualize import matshow3d

def plot_img_operations(pair, every_n=10):
    if isinstance(pair, list):
        print(f"metadata of patch idx 0 out of {len(pair)} samples.\n")
        p_0 = pair[0]
    else:
        p_0 = pair
    print("\n")
    p_0 = p_0["image"]
    if len(p_0.shape) > 4:
        p_0 = monai.utils.first(p_0)
    print(f"pixdim: {p_0.pixdim}")
    print(f"shape: {p_0.shape}")
    plt.figure(figsize=(15, 8))
    plt.subplot(1, 2, 1)
    matshow3d(pair["image"], every_n=every_n, fig=plt.gca())
    plt.subplot(1, 2, 2)
    matshow3d(pair["label"], every_n=every_n, fig=plt.gca())
    plt.show()

images = [
    {
        "image": "/workspace/Data/Task09_Spleen/imagesTr/spleen_10.nii.gz",
        "label": "/workspace/Data/Task09_Spleen/labelsTr/spleen_10.nii.gz"
    },
    {
        "image": "/workspace/Data/Task09_Spleen/imagesTr/spleen_8.nii.gz",
        "label": "/workspace/Data/Task09_Spleen/labelsTr/spleen_8.nii.gz"
    }
]
# image patch sampler
n_samples = 5
sampler = Compose([
    LoadImaged(keys=["image", "label"], ensure_channel_first=True),
    OrientationD(keys=("image", "label"), axcodes="SPL"),
    RandCropByPosNegLabeld(keys=["image", "label"], label_key = "label", spatial_size=(25, 256, 256), num_samples=n_samples, pos=1, neg=0.01)
])
# patch-level intensity shifts
patch_intensity = RandShiftIntensityd(keys="image", offsets=1.0, prob=1.0)
# construct the patch dataset
ds = PatchDataset(data=images,
                    patch_func=sampler,
                    samples_per_image=n_samples,
                    transform=patch_intensity)

train_patches_dataset = CacheDataset(data=ds, cache_rate=1.0, num_workers = 5) 

for item in DataLoader(train_patches_dataset, batch_size=2, shuffle=True, num_workers=2):
    plot_img_operations(item)
    break

Could you please test it on your dataset or share a piece if code so that I can reproduce the issue?
Hope it can help you, thanks!

[fogunsan]

Hi @KumoLiu, here is the code I am testing on:

import numpy as np

import monai
import matplotlib.pyplot as plt
from monai.data import PatchDataset, DataLoader, CacheDataset
from monai.transforms import RandCropByPosNegLabeld, RandShiftIntensityd, Compose, LoadImaged, OrientationD
from monai.visualize import matshow3d

images = [
    {
        "image": sorted(glob('*resampled_T1w.nii.gz'))[0],
        "label": sorted(glob('*normalized_fcm.nii.gz'))[0]
    },
    {
        "image": sorted(glob('*resampled_T1w.nii.gz'))[1],
        "label": sorted(glob('*normalized_fcm.nii.gz'))[1]
    }
]
# image patch sampler
n_samples = 20
sampler = Compose([
    LoadImaged(keys=["image", "label"], ensure_channel_first=True),
    OrientationD(keys=("image", "label"), axcodes="SPL"),
    RandCropByPosNegLabeld(keys=["image", "label"], label_key = "label", spatial_size=(64,64,64), num_samples=n_samples, pos=1, neg=0.01)
])
# patch-level intensity shifts
patch_transforms = Compose([Rand3DElasticd(keys =["image", "label"], sigma_range = (0.5,1), magnitude_range = (0.1, 0.3), prob=0.2, shear_range=(0.1, -0.05, 0.0, 0.0, 0.0, 0.0), scale_range=0.5, padding_mode= "zeros"),
          RandFlipd(keys =["image", "label"], prob = 0.4, spatial_axis=1)])

# construct the patch dataset
ds = PatchDataset(data=images,
                    patch_func=sampler,
                    samples_per_image=n_samples,
                    transform=patch_transforms)

train_patches_dataset = CacheDataset(data=ds, cache_rate=1.0, num_workers = 1) 

for item in DataLoader(train_patches_dataset, batch_size=2, shuffle=True, num_workers=1):
    images = item['image'].cuda()# gad images of batch
    labels = item['label'].cuda() # nongad images of batch
    fig, axes = plt.subplots(1,2)
    axes[0].imshow(images[0,0,5,:,:].cpu().data.numpy(), cmap="gray")
    axes[1].imshow(labels[0,0,5,:,:].cpu().data.numpy(), cmap="gray")

The first screenshot I took was when I set num_workers to 1. All the output patches produced were well aligned. The second screenshot was when I set num_workers to 8. There are a couple patches that are misaligned. The two images were rigidly registered and QC'ed.

[KumoLiu]

Hi @fogunsan, I just get your points, I don't think we can pass a dataset to CacheDataset since CacheDataset can only cache non-random transforms. And also most of the random transforms are non-thread safe.

MONAI/monai/data/dataset.py

Lines 735 to 740 in 98b6a15

	Note:
	`CacheDataset` executes non-random transforms and prepares cache content in the main process before
	the first epoch, then all the subprocesses of DataLoader will read the same cache content in the main process
	during training. it may take a long time to prepare cache content according to the size of expected cache data.
	So to debug or verify the program before real training, users can set `cache_rate=0.0` or `cache_num=0` to
	temporarily skip caching.

There is no way to even determine which ones are random transorm because transform is just None type.

MONAI/monai/data/dataset.py

Lines 804 to 805 in 98b6a15

	if not isinstance(transform, Compose):
	transform = Compose(transform)

Hope it can help you!
Thanks!

[fogunsan]

Because RandCropByPosNegLabeld is a random transform, does that mean patches from paired data cannot be reliably cached? Is there a way you'd suggest getting around this issue when dealing with a large dataset to still maximize loading and training speeds? I played around with multiprocessing with the dataloader but with the patches not cached, the overall training times are significantly reduced compared to when caching is implemented.

[KumoLiu]

Hi @fogunsan, I don't think you need to use PatchDataset. You can just put all the transforms in one Compose, then use CacheDataset so that all the non-random transforms can be cached, it will also accelerate your training. Also, you can add ToDeviced to move data to GPU and cache with CacheDataset, then execute random transforms on GPU directly, avoiding CPU -> GPU sync in every epoch.
For more details, you can refer to these two tutorials.
https://github.com/Project-MONAI/tutorials/blob/main/acceleration/dataset_type_performance.ipynb
https://github.com/Project-MONAI/tutorials/blob/main/acceleration/fast_training_tutorial.ipynb

Hope it can help you, thanks!

[fogunsan]

Okay, that makes sense thank you! It seems like reading from files that already are patched up, caching a portion of them with a cachedataset and then applying randomizable transforms on the fly with a threaddataloader during training seems like the most efficient route to take when dealing with a large dataset of patches. I had a final question about saving patches and whether MONAI has a function to read an image and label and save smaller patches into an output file (preserving correspondence between patches), similar to convert3d's export patches function. The only way I would think would be to iterate through batches in a Dataset and save the patches using monai.utils.save_image but this does not seem to allow for appending to a file that already exists (to generate training and validation patch files) or the ability to save paired data.

Thank you for all your help!

[KumoLiu]

Hi @fogunsan, maybe you can use matshow3d.

MONAI/monai/visualize/utils.py

Lines 97 to 100 in e6ec945

	>>> volumes = [np.random.rand(1, 10, 10, 10), np.random.rand(1, 10, 10, 10)]
	>>> fig = plt.figure()
	>>> matshow3d(volumes, fig=fig, title="List of Volumes")
	>>> plt.show()