Lazy DivisiblePadd converts uint8 to float32

[john-zielke-snkeos]

Describe the bug
When running DivisiblePadd in a lazy Compose pipeline, the padding implicitly converts a uint8 tensor to a float32 tensor.
This happens only with lazy=True, if I run the same pipeline with lazy=False, the tensor stays uint8 type.

To Reproduce

import torch
from monai.transforms import DivisiblePadd, Compose

# create a 3D tensor
input_batch = {"x": torch.arange(0, 27,dtype=torch.uint8).reshape(3, 3, 3)}
assert input_batch["x"].dtype == torch.uint8

print("non lazy transform")
non_lazy_transforms = Compose([DivisiblePadd(k=5,keys=["x"])],lazy=False)
non_lazy_batch = non_lazy_transforms(input_batch)
print(non_lazy_batch["x"].dtype)
assert non_lazy_batch["x"].dtype == torch.uint8

print("lazy transform")
lazy_transforms = Compose([DivisiblePadd(k=5,keys=["x"])],lazy=True)
lazy_batch = lazy_transforms(input_batch)
print(lazy_batch["x"].dtype)
assert lazy_batch["x"].dtype == torch.uint8

Output:

non lazy transform
torch.uint8
lazy transform
torch.float32
Traceback (most recent call last):
  File "/mnt/c/Projects/monai_bug/padd_bug.py", line 18, in <module>
    assert lazy_batch["x"].dtype == torch.uint8
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

Expected behavior
The type does not get converted, as this causes problems later on, especially for label maps.

Environment

Ensuring you use the relevant python executable, please paste the output of:

================================
Printing MONAI config...

MONAI version: 1.2.0
Numpy version: 1.25.0
Pytorch version: 2.0.1
MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False
MONAI rev id: c33f1ba

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
ITK version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: NOT INSTALLED or UNKNOWN VERSION.
scikit-image version: NOT INSTALLED or UNKNOWN VERSION.
Pillow version: 9.4.0
Tensorboard version: NOT INSTALLED or UNKNOWN VERSION.
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.15.2
tqdm version: NOT INSTALLED or UNKNOWN VERSION.
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: NOT INSTALLED or UNKNOWN VERSION.
pandas version: NOT INSTALLED or UNKNOWN VERSION.
einops version: NOT INSTALLED or UNKNOWN VERSION.
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
pynrrd version: NOT INSTALLED or UNKNOWN VERSION.

[wyli]

Thanks for reporting, it is a known issue (#6340), the root cause is a limitation from torch.nn.functional.grid_sampler which only supports floats. At the moment float32 is always used in the lazy resampling even if some specific cases are achievable by simple array slicing without changing the data types. This is to ensure the preprocessing collate can work robustly regardless of those specific cases.

please see also

• Spatial transforms / RandAffined/ Resampler should not change data type #5511
• Rotate with nearest filter does not return same dtype #3831

[john-zielke-snkeos]

Thank you for the pointer, I did not know this is a general limitation of the spatial transforms right now. Since there seems to be no ETA for the needed changes in pytorch yet, is there any plan to include a check in Monai to stop the lazy transform or at least print a warning in case this functionality is triggered?

[wyli]

agreed @john-zielke-snkeos , marking this as a feature request.

[atbenmurray]

@KumoLiu We should keep track of this issue; I think we'll be fixing it with a more comprehensive consolidated resampling function.