Mixing MONAI and TorchIO transforms in OneOf with batched data leads to RuntimeError

Describe the bug Mixing MONAI and TorchIO transforms in the same OneOf instance with batched data leads to RuntimeError: each element in list of batch should be of equal size, even though both transforms within OneOf seem to produce tensors of the same size. Using MONAI transforms only or TorchIO transforms only within a given OneOf instance works with batched data. Would like to be able to combine both within a given OneOf instance.

To Reproduce The following example, adapted from the 3D spleen segmentation tutorial, should reproduce the error:

import glob
import os
import tempfile

from monai.apps import download_and_extract
from monai.config import print_config
from monai.data import DataLoader, Dataset
import monai.transforms as tf
from monai.utils import first, set_determinism
import torchio as tio

print_config()

directory = os.environ.get("MONAI_DATA_DIRECTORY")
root_dir = tempfile.mkdtemp() if directory is None else directory
print(root_dir)

resource = "https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar"
md5 = "410d4a301da4e5b2f6f86ec3ddba524e"

compressed_file = os.path.join(root_dir, "Task09_Spleen.tar")
data_dir = os.path.join(root_dir, "Task09_Spleen")
if not os.path.exists(data_dir):
    download_and_extract(resource, compressed_file, root_dir, md5)

train_images = sorted(
    glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
train_labels = sorted(
    glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
data_dicts = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]
train_files, val_files = data_dicts[:-9], data_dicts[-9:]

set_determinism(seed=0)

oneof_transforms = [
    tf.RandAdjustContrastd(
        keys=['image'],
        prob=1.0
    ),
    tf.Compose([
        tf.ToTensord(keys=['image', 'label']),
        tio.transforms.RandomBlur(
            p=1.0,
            include=['image']
        ),
        tf.ToNumpyd(keys=['image', 'label'])
    ])
]

train_transforms = tf.Compose([
    tf.LoadImaged(keys=["image", "label"]),
    tf.EnsureChannelFirstd(keys=["image", "label"]),
    tf.Spacingd(keys=["image", "label"], pixdim=(
        1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
    tf.Orientationd(keys=["image", "label"], axcodes="RAS"),
    tf.ScaleIntensityRanged(
        keys=["image"], a_min=-57, a_max=164,
        b_min=0.0, b_max=1.0, clip=True,
    ),
    tf.CropForegroundd(keys=["image", "label"], source_key="image"),
    tf.OneOf(transforms=oneof_transforms),
    tf.RandCropByPosNegLabeld(
        keys=["image", "label"],
        label_key="label",
        spatial_size=(96, 96, 96),
        pos=1,
        neg=1,
        num_samples=4,
        image_key="image",
        image_threshold=0,
    ),
    tf.EnsureTyped(keys=["image", "label"]),
])

ds = Dataset(data=train_files, transform=train_transforms)
loader = DataLoader(ds, batch_size=2, shuffle=True)

for batch_data in loader:
    continue

Expected behavior I would expect the dataloader to be able to collate the samples within a batch if the two samples have the same size (e.g., [4, 1, 96, 96, 96] for the example above) even if they’re transformed by transforms from different libraries (i.e., MONAI and TorchIO). With that said, I’m not too familiar with the low-level details of the dataloader and collate functions, so I could be missing an important detail.

Screenshots

Environment

================================
Printing MONAI config...
================================
MONAI version: 0.7.dev2133
Numpy version: 1.19.2
Pytorch version: 1.8.1+cu102
MONAI flags: HAS_EXT = False, USE_COMPILED = False
MONAI rev id: 20ffa3f987fad60a8428ec635fb0b4f6ccca9747

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: 3.2.1
scikit-image version: NOT INSTALLED or UNKNOWN VERSION.
Pillow version: 7.2.0
Tensorboard version: NOT INSTALLED or UNKNOWN VERSION.
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.9.1+cu102
tqdm version: 4.60.0
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: NOT INSTALLED or UNKNOWN VERSION.
pandas version: 1.2.4
einops version: NOT INSTALLED or UNKNOWN VERSION.

For details about installing the optional dependencies, please visit:
    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies


================================
Printing system config...
================================
`psutil` required for `print_system_info`

================================
Printing GPU config...
================================
Num GPUs: 0
Has CUDA: False
cuDNN enabled: True
cuDNN version: 7605

Additional context N/A

[rijobro]

I’ve done some messing around and think the problem is due to differences in the image_transforms. If RandAdjustConstrastd is used, an extra element will be in the list compared to the torchio transform. Adding tf.DeleteItemsd(keys=["image_transforms", "label_transforms"]), at the end of the transforms worked for me.

But certainly something that we need to think about for avoiding this in the future. Thanks for raising this!

[lyndonboone]

Thanks @rijobro , that seems to do the trick! Much appreciated.