Meta_dict and lazy resampling in apply transform

[tangy5]

Hi Team, I'm not sure if this is a bug or it's expected.

When I set env var USE_META_DICT=1, some transforms will fail due to the "lazy" argument. If USE_META_DICT=0 or none, it works normally.

Currently, this is causing some functions to fail when USE_META_DICT is set to true (1).

Here is a code snipet for reproducing:

1. run export USE_META_DICT=1
2. Create a python script with following codes lazy_meta_dict.py, then run python lazy_meta_dict.py

from functools import partial

import numpy as np
import torch
from monai.data import CacheDataset, create_test_image_3d
from monai.transforms import (
    AddChanneld,
    Compose,
    CropForegroundd,
    DivisiblePadd,
)

def get_data(num_examples, input_size, data_type=np.asarray, include_label=True):
    custom_create_test_image_3d = partial(
        create_test_image_3d, *input_size, rad_max=7, num_seg_classes=1, num_objs=1
    )
    data = []
    for _ in range(num_examples):
        im, label = custom_create_test_image_3d()
        d = {}
        d["image"] = data_type(im)
        d["image_meta_dict"] = {"affine": np.eye(4)}
        if include_label:
            d["label"] = data_type(label)
            d["label_meta_dict"] = {"affine": np.eye(4)}
        d["label_transforms"] = []
        data.append(d)
    return data[0] if num_examples == 1 else data

def test_epistemic_scoring():
    input_size = (20, 20, 20)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    keys = ["image", "label"]
    num_training_ims = 10
    train_data = get_data(num_training_ims, input_size)
    # print("Hey!!:{}".format(CropForegroundd.__dict__.items()))
    transforms = Compose(
        [
            AddChanneld(keys),
            CropForegroundd(keys, source_key="image"),
            DivisiblePadd(keys, 4),
        ]
    )

    train_ds = CacheDataset(train_data, transforms)

if __name__ == "__main__":
    test_epistemic_scoring()

Expected error:

  File "/home/yucheng/anaconda3/lib/python3.9/site-packages/monai/transforms/transform.py", line 145, in apply_transform
    return _apply_transform(transform, data, unpack_items, lazy, overrides, log_stats)
  File "/home/yucheng/anaconda3/lib/python3.9/site-packages/monai/transforms/transform.py", line 102, in _apply_transform
    return transform(data, lazy=lazy) if isinstance(transform, LazyTrait) else transform(data)
TypeError: wrapper() got an unexpected keyword argument 'lazy'

[KumoLiu]

Seems we missed passing lazy in the wrapper.

MONAI/monai/transforms/transform.py

Lines 389 to 391 in e6ec945

	if config.USE_META_DICT:
	# call_update after MapTransform.__call__
	cls.__call__ = transforms.attach_hook(cls.__call__, MapTransform.call_update, "post") # type: ignore

MONAI/monai/transforms/utils.py

Line 1753 in e6ec945

def wrapper(inst, data):

[wyli]

Thanks, I can replicate this issue in 1.2.0 (but not 1.2.0rc7), it's introduced by #6537. The root cause is that the API changes to the __call__ of MapTransform subclasses in the lazy resampling refactoring, for example in monai/transforms/spatial/dictionary.py:

they are inconsistent with the base class MapTransform data assumptions that we've been using all the time:

MONAI/monai/transforms/transform.py

Lines 430 to 445 in e6ec945

	def __call__(self, data):
	"""
	``data`` often comes from an iteration over an iterable,
	such as :py:class:`torch.utils.data.Dataset`.
	To simplify the input validations, this method assumes:
	- ``data`` is a Python dictionary,
	- ``data[key]`` is a Numpy ndarray, PyTorch Tensor or string, where ``key`` is an element
	of ``self.keys``, the data shape can be:
	#. string data without shape, `LoadImaged` transform expects file paths,
	#. most of the pre-/post-processing transforms expect: ``(num_channels, spatial_dim_1[, spatial_dim_2, ...])``,
	except for example: `AddChanneld` expects (spatial_dim_1[, spatial_dim_2, ...])
	- the channel dimension is often not omitted even if number of channels is one.

@Nic-Ma @ericspod we may have to release 1.2.1 to fix these. Let's discuss this in today's dev meeting. cc @mmodat

[atbenmurray]

@tangy5 please reach out to me so I can further understand the problem. I'm fresh to this as of the meeting

[atbenmurray]

@tangy5 let's schedule a quick chat on monday, so we can figure out the best way to resolve it, if you are free.

[tangy5]

@tangy5 let's schedule a quick chat on monday, so we can figure out the best way to resolve it, if you are free.

Thank you so much. Currently, this is not a blocker issue for other MONAI related platforms such as MONAI Label and MONAI Toolkit, but we'd better discuss whether we will make it compatible with meta_dict. If we want to deprecate the meta_dict, we can tell users to use metaTensor instead, if users are using meta_dict with moani=1.2.0, we can tell them to shift to metaTensor or downgrade monai.

In both cases, it should work. Thank you! I'm available for a chat on Monday if a discussion is needed.

[Nic-Ma]

@atbenmurray ,

Please feel free to organize the discussion meeting.
I think @wyli 's suggestion is one of the proper solutions, maybe we should keep the data: Dict input simple and make lazy as one item in the data, similar to how we add new key: value in the data: https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/io/dictionary.py#L178.
CC @ericspod @wyli for more discussion.

Thanks.

[atbenmurray]

Hi @Nic-Ma I'd like to look for other options as well.

It seems to be very restrictive for dictionary based transforms to be forced to only accept one call parameter. It isn't true for non-dictionary transforms, and makes us have to have two different mechanisms for essentially the same piece of functionality. It's almost always better to have one mechanism instead of two where possible.

I'm trying to understand why it isn't simply a case of extending wrapper to take args and kwargs. That is a pretty normal pattern for such a function. See torch internals for many examples of using args and kwargs to forwarg general arguments to a call site after doing their own behind the scenes work.

[atbenmurray]

I've put together a draft PR #6598 that solves the problem. It might need tweaking, of course.

[Nic-Ma]

Hi @atbenmurray ,

I think the initial reason for accepting only 1 call parameter in the dict transforms is that:
We designed to chain the dict transform and didn't plan to chain the array transforms, to ensure dict transforms are consistent with the same input and output in a chain, 1 dict data is more clear.
CC @wyli @ericspod to add more if I missed anything.

Thanks.

[atbenmurray]

It would make more sense if we also applied the same restriction to array transforms, but we don't. Allowing additional arguments for array transforms but not for dictionary transforms means that we need different code paths elsewhere depending on whether we are handling array transforms and dictionary transforms. In general, I think we should minimize discrepancies if we can.

Have you had a look at the PR? It is a really simple and small fix and seems to do away with the need to restrict dictionary call in this way.

[wyli]

I think the initial reason for accepting only 1 call parameter in the dict transforms is that: We designed to chain the dict transform and didn't plan to chain the array transforms, to ensure dict transforms are consistent with the same input and output in a chain, 1 dict data is more clear. CC @wyli @ericspod to add more if I missed anything.

Thanks.

yes, another benefit is that this also allows for the subclass implmentations to dynamically determine the additional arguments based on the value of specific data keys.