Handling randomness with multiple workers using ZipDataset

[maxwellreynolds]

Hello,

I am having some difficulty with monai random transformations. When I train over multiple epochs, the same images are produced (i.e. the augmentations are not random).

I am using monai datasets and dataloaders (also applies if I use a torch dataloader instead). My transforms are:

train_transforms = Compose([ScaleIntensity(), RandGaussianSmooth(sigma_x=(0.0,1.5),sigma_y=(0.0,1.5),sigma_z=(0.0,1.5), prob=1.0),
AddChannel(),
Resize([121,145,121]),
RandSpatialCrop((96, 96, 96), random_center = True, random_size=False),
EnsureType()])

This issue does not apply to torch random transforms. For example, when I add a torch transform such as RandomAffine, the images are not constant across epochs.

It also does not apply when using num_workers = 0

I have tried passing in a worker_init_fn (from the faq) but still getting the same problem.

    def worker_init_fn_faq(worker_id):
        worker_info = torch.utils.data.get_worker_info()
        try:
            worker_info.dataset.transform.set_random_state(worker_info.seed % (2 ** 32))
        except AttributeError:
            pass

I think the problem may be related to this issue

Appreciate any help.

[maxwellreynolds]

I guess the way to do this here (without shuffling) is call randomize() on the dataloader's ImageDataset, in my case train_loader.dataset.data[0].randomize() , after each epoch

[wyli]

could you please provide a small example to show the unexpected results? perhaps you can try the default work_init_fn

MONAI/monai/data/utils.py

Line 655 in 32af32c

def worker_init_fn(worker_id: int) -> None:

[maxwellreynolds]

@wyli thank you for the response. I think the issue is actually with ZipDataset (not having Randomizable?). Using the default worker_init_fn does not change the results, but using just ImageDataset resolves the issue. Here is a reproducible version:

from monai.transforms import RandGaussianSmooth,  AddChannel, RandSpatialCrop, Compose, ScaleIntensity, EnsureType
import pdb
import monai
from monai.data import ImageDataset, Dataset, ArrayDataset, ZipDataset, DataLoader
import scipy
import numpy as np
import torch
import torchvision.transforms as transforms
import nibabel as nib
from monai.data.utils import worker_init_fn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def main():
    print('device',device)
    #create 4 (3x3) test images
    test_img1=torch.Tensor([[1.,2.,3.],[4.,5.,6.,],[7.,8.,9.]])
    test_img2=test_img1*2.
    test_img3=test_img1*3.
    test_img4=test_img1*4.
    test_img1=nib.Nifti1Image(test_img1.numpy(),None)
    test_img2=nib.Nifti1Image(test_img2.numpy(),None)
    test_img3=nib.Nifti1Image(test_img3.numpy(),None)
    test_img4=nib.Nifti1Image(test_img4.numpy(),None)
    #write images
    nib.save(test_img1, 'test_img1.nii.gz')
    nib.save(test_img2, 'test_img2.nii.gz')
    nib.save(test_img1, 'test_img3.nii.gz')
    nib.save(test_img2, 'test_img4.nii.gz')
    ds_transforms = Compose([AddChannel(),
    RandGaussianSmooth(sigma_x=(0,1),sigma_y=(0,1),sigma_z=(0,1), prob=1),
    # RandRotate((-180,180),prob=1.),
    RandSpatialCrop((2,2), random_center = True, random_size=False),
    EnsureType(),
    ])
    ds1=ImageDataset(image_files = ['test_img1.nii.gz','test_img2.nii.gz','test_img3.nii.gz','test_img4.nii.gz',],
    labels=[1,2,3,4,],
    transform = ds_transforms, 
    reader="NibabelReader", 
    )
    ds2 = Dataset(data = [1,2,3,4,])
    ds=ZipDataset([ds1,ds2])
    ds_loader = DataLoader(ds, batch_size = 2, num_workers = 4,
    worker_init_fn = worker_init_fn,
    )
    sigma_min,sigma_max=0,1.5
    prob_blur = 1
    for epoch in range(0,5):
        print('Epoch [{0}]: \t'.format(epoch))
        train(ds_loader)

def train(train_loader):
    for i, (input, target, another_label) in enumerate(train_loader):
        if i==0:
            image_0 = input[0,0,:,:].detach().cpu().numpy()
            print('image0:',image_0)
    return None
if __name__ == '__main__':
    main()

Output:
device cuda:0 Epoch [0]: image0: [[3.9185538 4.9889097] [5.3457174 6.3036075]] Epoch [1]: image0: [[3.9185538 4.9889097] [5.3457174 6.3036075]] Epoch [2]: image0: [[3.9185538 4.9889097] [5.3457174 6.3036075]] Epoch [3]: image0: [[3.9185538 4.9889097] [5.3457174 6.3036075]] Epoch [4]: image0: [[3.9185538 4.9889097] [5.3457174 6.3036075]]

[wyli]

many thanks, indeed that's a bug for zipdataset, I'll make a patch. -- patch merged #4991