RuntimeError with SpatialPad when input is MetaTensor and padding size is the same as tensor size

[virginiafdez]

Hello,

When dealing with a MetaTensor (track_meta = True) that has the same shape you are padding to with SpatialPad, the following error occurs:

Traceback (most recent call last):
  File "/home/vf19/.local/share/JetBrains/Toolbox/apps/PyCharm-P/ch-0/202.7660.27/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "/home/vf19/PycharmProjects/semantic_conditioning/venv/lib/python3.9/site-packages/monai/transforms/croppad/array.py", line 168, in __call__
    return pad_func(img_t, to_pad_, self.get_transform_info(), mode_, lazy_, **kwargs_)
  File "/home/vf19/PycharmProjects/semantic_conditioning/venv/lib/python3.9/site-packages/monai/transforms/croppad/functional.py", line 196, in pad_func
    meta_info = TraceableTransform.track_transform_meta(
  File "/home/vf19/PycharmProjects/semantic_conditioning/venv/lib/python3.9/site-packages/monai/transforms/inverse.py", line 188, in track_transform_meta
    affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (72x3 and 24x24)

This doesn't happen if input is a tensor, or if I set track_meta to False with a transform prior to padding.

To Reproduce

1. Use SpatialPad, not SpatialPadd
2. Input a MetaTensor with track_meta = True
3. input tensor has shape: [24, 8, 24, 32], SpatialPad gets passed spatial_size = [-1, 24, 32]

Expected behavior

If this specific scenario, where the shape should be different, is not supported in case of MetaTensors, the error should be clearer. In any other scenario (dealing with normal tensors) the shape is left the same and no error occurs.

Environment

python 3.9
monai 1.3.0
pytorch 2.1.0, 1.13.1 (both have the same problem)

[KumoLiu]

Hi @virginiafdez, thanks for the reporting.
I can't reproduce the issue, could you please try the code below?

import torch
from monai.transforms import SpatialPad
from monai.data import MetaTensor

data = MetaTensor(torch.rand([24, 8, 24, 32]), track_meta=True)

out = SpatialPad([-1, 24, 32])(data)

Thanks!

[ericspod]

This issue is related to how data is being collated by the data loader when putting list-based batches together with RandSpatialCropSamplesd. We're still trying to narrow down what's going on but it's more than just the transform that's raising the error.

[virginiafdez]

The following replicates the issue:

import monai.transforms as mt
import torch
from monai.data.dataloader import DataLoader
from monai.data.dataset import Dataset
input_dict = [{'im': torch.randn([1, 192, 256, 256])}, {'im': torch.randn([1, 192, 256, 256])}, {'im': torch.randn([1, 192, 256, 256])},
              {'im': torch.randn([1, 192, 256, 256])}, {'im': torch.randn([1, 192, 256, 256])}, {'im': torch.randn([1, 192, 256, 256])}]
transforms = mt.Compose(
     [mt.RandSpatialCropSamplesd(keys=['im'],
                                num_samples=8,
                                max_roi_size=[192, 256, 1],
                                roi_size=[192, 256, 1],
                                random_size=False),
     mt.SqueezeDimd(keys=['im'], dim=-1),
     mt.ToTensord(keys=['im'])])
dataset = Dataset(input_dict, transforms)
element = next(iter(dataset))
dataloader = DataLoader(dataset, batch_size=2, num_workers=0)
post_transform = mt.SpatialPad(spatial_size=[192, 256])
input = torch.randn([192, 256])
for b_el in dataloader:
    image = b_el['im']
    print(image.shape,image.affine.shape)
    image_tf = post_transform(image[:,0])

After cropping the samples and putting these transform into the loader, there seems to be a problem with len(orig_affine) in TraceableTransform.track_transform_meta.

@atbenmurray in case you want to have a look / help!

[ericspod]

One thing to note in the above code is that post_transform is being applied to a whole batch of images rather than individually so the dimension of the input isn't what's expected. If the last line is image_tf = post_transform(image[0]) this example will work. The issue with the matrix multiplication is because of this batching actually but there should be some check to indicate that this is the case rather than allowing an unhelpful exception like this.

[KumoLiu]

Just like Eric said, the issue is due to that the post_transform is applied on a batch of data.
One more thing you could do is that using decollate_batch.

image_tf = [post_transform(i) for i in decollate_batch(image)]

Also create a patch to add more information when raising exception.

[vikashg]

closing as merged @KumoLiu @virginiafdez