ArrayDataset only applies same transform if transforms match

[rijobro]

Describe the bug
Noticed when looking into #2052.

ArrayDataset gives matching results if img_transform and label_transform match each other. However results diverge when they don't match. At which point, I don't fully understand the utility of the ArrayDataset over using dictionary transforms. Perhaps I've misunderstood something.

To Reproduce
The following code is fine if t1 matches t2, but fails if they are different.

works:

t1 = RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)
t2 = RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)

doesn't work

t1 = Compose([RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)])
t2 = RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)

also doesn't work

t1 = Compose([Lambda(lambda x: x), RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)])
t2 = RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)

import sys
import unittest
import numpy as np

from monai.data import ArrayDataset, DataLoader
from monai.transforms import Compose, RandSpatialCropSamples


class TestArrayDatasetTransforms(unittest.TestCase):
    def test_same_transforms(self):

        im = np.arange(0, 100 ** 3).reshape(1, 100, 100, 100)
        a1 = [np.copy(im) for _ in range(20)]
        a2 = [np.copy(im) for _ in range(20)]
        num_samples = 10
        t1 = Compose([RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)])
        t2 = RandSpatialCropSamples(roi_size=(10, 10, 10), num_samples=num_samples, random_size=False)

        dataset = ArrayDataset(a1, t1, a2, t2)
        self.assertEqual(len(dataset), len(a1))
        dataset.set_random_state(1234)
        n_workers = 0 if sys.platform != "linux" else 2
        batch_size = 2
        loader = DataLoader(dataset, batch_size=batch_size, num_workers=n_workers)
        batch_data = next(iter(loader))
        self.assertEqual(len(batch_data), num_samples * 2)
        self.assertEqual(batch_data[0].shape[0], batch_size)
        for i in range(num_samples):
            out1, out2 = batch_data[i], batch_data[i + num_samples]
            np.testing.assert_array_equal(out1, out2)


if __name__ == "__main__":
    unittest.main()

[wyli]

that's because Compose's __call__ is implemented with this logic:

MONAI/monai/transforms/transform.py

Lines 46 to 48 in 299f2c2

	if isinstance(data, (list, tuple)) and map_items:
	return [transform(item) for item in data]
	return transform(data)

the idea is that for example in

Compose([RandSpatialCropSamples(num_samples=10, ...), OtherTransform(...)]

the OtherTransform will be applied to each of the 10 samples generated by the random cropping respectively (instead of applying OtherTransform to the list of 10 samples). perhaps we can add an option to switch between these two modes, I'll make a PR...

[rijobro]

I don't think this is related to the number of samples. The same behaviour can be seen using any other random transform, e.g., RandAffine:

import unittest
import numpy as np

from monai.data import ArrayDataset, DataLoader
from monai.transforms import Compose, RandAffine


class TestArrayDatasetTransforms(unittest.TestCase):
    def test_same_transforms(self):

        im = np.arange(0, 100 ** 3).reshape(1, 100, 100, 100)
        a1 = [np.copy(im) for _ in range(20)]
        a2 = [np.copy(im) for _ in range(20)]

        t1 = Compose([RandAffine(prob=1, rotate_range=(np.pi, np.pi),
                                 translate_range=(5, 5), as_tensor_output=False)])
        t2 = RandAffine(prob=1, rotate_range=(np.pi, np.pi),
                        translate_range=(5, 5), as_tensor_output=False)

        dataset = ArrayDataset(a1, t1, a2, t2)
        dataset.set_random_state(1234)
        loader = DataLoader(dataset, batch_size=2, num_workers=0)
        batch_data = next(iter(loader))
        np.testing.assert_array_equal(batch_data[0], batch_data[1])


if __name__ == "__main__":
    unittest.main()

[wyli]

I see, the set random state method of compose adds some control of the transforms

MONAI/monai/transforms/compose.py

Lines 114 to 118 in a6ca680

	super().set_random_state(seed=seed, state=state)
	for _transform in self.transforms:
	if not isinstance(_transform, Randomizable):
	continue
	_transform.set_random_state(seed=self.R.randint(MAX_SEED, dtype="uint32"))

I feel that is somehow expected, to have a sequence of random seeds for a list of transforms.