Make `apply_pending_transforms` optional in `execute_compose`

[function2-llx]

Is your feature request related to a problem? Please describe. Suppose I'm defining transform as follows:

from monai import transforms as mt
crop_t = mt.OneOf(
    [
        mt.RandSpatialCropD(['img', 'seg'], roi_size=(8, 8)),
        mt.RandCropByPosNegLabelD(['img', 'seg'], label_key='seg', spatial_size=(8, 8)),
    ],
    lazy=True,
)
rotate_t = mt.RotateD(['img', 'seg'], np.pi / 3)
t = mt.Compose([crop_t, rotate_t], lazy=True)

What I intend to do here is to crop a patch from the image with one of the two methods (similar to what nnU-Net does), and rotate the cropped patch afterwards. It will be very nice if I can lazy resample the patch after both crop and rotation (, and which should be possible). However, OneOf is a subclass of Compose, and it calls execute_compose, which always calls apply_pending_transforms at last, which makes the cropped patch resampled in advance and causes information loss during the subsequent rotation.

https://github.com/Project-MONAI/MONAI/blob/0a3cdd9f1403cbad057b96230840ee17cf430dbf/monai/transforms/compose.py#L476-L487

https://github.com/Project-MONAI/MONAI/blob/0a3cdd9f1403cbad057b96230840ee17cf430dbf/monai/transforms/compose.py#L114

Describe the solution you'd like We can add a parameter to execute_compose, and the subclasses of Compose if necessary, with a name like apply_pendding indicating if the pending operations should be applied.

Additional context This was originally proposed in a discussion (the second question). It seems that my feature request is not clearly conveyed. But I believe this is a very useful feature (as my example shows), therefore I'd like to propose again here.

[wyli]

yes, this makes sense. also RandCropByPosNegLabelD needs label_key='seg' to be up to date but can be lazily applied to 'img', this is not currently implemented.

[function2-llx]

RandCropByPosNegLabelD needs label_key='seg' to be up to date

I agree. Fortunately, such crop operation is usually performed in early stage, e.g., for nnU-Net, after resampling to a target spacing (which is done during pre-processing).

[atbenmurray]

Hi @function2-llx.

There are two different aspects to this problem.

The first is that we have a number of limitations in the current Compose implementation when using nested Compose and subclasses that this particular issue can feed into. There is issue #7130 and PR #7140 which is focussed on handling nested Compose, and is probably the correct place to handle the apply pending situation.

The second issue is more related to the capabilities of lazy resampling. I'm hoping that we can extend the functionality in 1.4 so that transforms such as RandCropByPosNegLabelD, that require up date data in order to execute do the following:

1. The transform copies and then internally resamples the data to be up to date, creating a temporary tensor for the results
2. The transform performs its calculations on the temporary tensor
3. The transform applies the pending operation to the un-resampled data and returns it,

This allows such transforms to execute on up to date data but not break the flow of pending transforms and is part of the roadmap for a full lazy resampling implementation.

[atbenmurray]

@wyli The first part of this issue relates to the problem with nested compose raised in #7130 . It makes sense to fix it as part of #7140, The second part is part of the lazy resampling feature set and has an entry in #7151

[function2-llx]

1. The transform copies and then internally resamples the data to be up to date, creating a temporary tensor for the results
2. The transform performs its calculations on the temporary tensor
3. The transform applies the pending operation to the un-resampled data and returns it,

This sounds great. I think it's also worth mentioning that the both options (1. creating a tensor for temporary usage and 2. applying the pending transforms permanently) perform the same set of resampling operations (i.e., resampling with the same grids). If the time complexity of resampling is only related to the grid size (it seems so), this method does not increase time complexity as well.

[atbenmurray]

1. The transform copies and then internally resamples the data to be up to date, creating a temporary tensor for the results
2. The transform performs its calculations on the temporary tensor
3. The transform applies the pending operation to the un-resampled data and returns it,

This sounds great. I think it's also worth mentioning that the both options (1. creating a tensor for temporary usage and 2. applying the pending transforms permanently) perform the same set of resampling operations (i.e., resampling with the same grids). If the time complexity of resampling is only related to the grid size (it seems so), this method does not increase time complexity as well.

With respect to grids, I'd like to get away from their use altogether if they aren't needed. For the most part, transforms that don't involve deformation really shouldn't be resampling with grids at all, as it is perfectly feasible to perform a GPU affine transform to 2D and 3D textures without generating a grid. I always felt like this was a curious omission from pytorch and it should certainly save time and memory

[atbenmurray]

This issue is fixed in PR #7140

[wyli]

This issue is fixed in PR #7140

Let's please not provide misleading information here. The PR #7140 has design issues (see the PR's comments) and multiple premerge testing errors, it's not very useful in the current form.

[atbenmurray]

This issue is fixed in PR #7140

Let's please not provide misleading information here. The PR #7140 has design issues (see the PR's comments) and multiple premerge testing errors, it's not very useful in the current form.

It is not misleading information. If you go to #7140 you see it is clearly a draft PR, but it fixes the issue. However, allow me to rephrase.

#7140 is a draft PR that fixes the issue described here.

[atbenmurray]

Taking this back on as part of PR #7140