Random transformation always do the same thing in every epoch

[jxchen01]

I am having a very strange issue, where all my random transformations produce exactly the same thing in every epoch. For example, I am using RandSpatialCropSamples to generate random samples in each epoch. But, it ends up that every epoch, the cropped samples are exactly the same. Does anyone has any idea of what could be wrong?

The code to reproduce this issue is attached below:

import pytorch_lightning
from monai.utils import set_determinism
from monai.transforms import (
    AddChanneld,
    Compose,
    CropForegroundd,
    LoadImaged,
    RandCropByPosNegLabeld,
    RandSpatialCropd,
    RandSpatialCropSamplesd,
    ScaleIntensityRanged,
    EnsureTyped,
    EnsureType,
)
from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.losses import DiceCELoss
from monai.data import PersistentDataset, list_data_collate

from typing import Dict, List, Sequence, Tuple, Union
import numpy as np
from monai.data import ImageReader
from monai.utils import ensure_tuple, require_pkg
from monai.config import PathLike
from monai.data.image_reader import _stack_images

from aicsimageio.writers import OmeTiffWriter
import torch
import tempfile
import os


def creat_random_data(num_images: int = 10):
    # Generate some image data
    tempdir = tempfile.mkdtemp()
    A_list = []
    B_list = []
    for img_idx in range(num_images):
        # image for domain A
        domainA_ = np.random.randint(0, 5, size=(64, 256, 128)).astype(np.uint32)

        # different size image for domain B
        domainB_ = np.random.randint(0, 2, size=(64, 256, 128)).astype(np.uint8)

        # save the images
        A_name = os.path.join(tempdir, "imgA_" + str(img_idx) + ".npy")
        B_name = os.path.join(tempdir, "imgB_" + str(img_idx) + ".npy")
        np.save(A_name, domainA_)
        np.save(B_name, domainB_)
        A_list.append(A_name)
        B_list.append(B_name)

    return A_list, B_list, tempdir


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

    def prepare_data(self):
        # create a random dataset
        train_A, train_B, tmpdir = creat_random_data(num_images=1)
        print(f"temporay directory created at {tmpdir}")

        train_files = [
            {"imageA": image_A, "imageB": label_B}
            for image_A, label_B in zip(train_A, train_B)
        ]

        # set deterministic training for reproducibility
        set_determinism(seed=0)

        # define the data transforms
        train_transforms = Compose(
            [
                LoadImaged(
                    keys=["imageA", "imageB"]
                ),
                AddChanneld(keys=["imageA", "imageB"]),
                RandSpatialCropSamplesd(
                    keys=["imageA", "imageB"],
                    random_size=False,
                    num_samples=1,
                    roi_size=(64, 64, 64),
                ),
                EnsureTyped(keys=["imageA", "imageB"]),
            ]
        )

        self.train_ds = PersistentDataset(
            data=train_files, cache_dir="./tmp/", transform=train_transforms
        )

    def train_dataloader(self):
        train_loader = torch.utils.data.DataLoader(
            self.train_ds,
            batch_size=1,
            num_workers=2,
            collate_fn=list_data_collate,
            persistent_workers=False,
        )
        return train_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),
            strides=(2, 2, 2),
            num_res_units=2,
            norm=Norm.BATCH,
        )
        self.loss_function = DiceCELoss(to_onehot_y=True, softmax=True)

    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):
        imageA, imageB = batch["imageA"], batch["imageB"]

        # check data in each batch
        OmeTiffWriter.save(imageA.detach().cpu().numpy(), str(self.current_epoch) + "_A.tiff", dim_order="TCZYX")
        OmeTiffWriter.save(imageB.detach().cpu().numpy(), str(self.current_epoch) + "_B.tiff", dim_order="TCZYX")

        output = self.forward(imageA)
        loss = self.loss_function(output, imageB)
        return {"loss": loss}


if __name__ == "__main__":

    # initialise the LightningModule
    net = Net()
    dm = myDataModule()
    dm.prepare_data()

    # initialise Lightning's trainer (basic pytorch lightning).
    trainer = pytorch_lightning.Trainer(
        gpus=1, precision=16, enable_checkpointing=True, logger=False, max_epochs=5
    )

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

[KumoLiu]

Hi @jxchen01, please use DataLoader in monai, it handles MONAI randomizable objects with appropriate random state
managements for deterministic behaviour. Or just update worker_init_fn in torch.utils.data.DataLoader with MONAI's. Hope it can helps you.

[jxchen01]

OMG... I have been using the wrong dataloader in the past few months .... Thanks so much for the great help @KumoLiu !!!!

[jxchen01]

it works perfectly now