Lazy resampling executes non-lazy transforms without executing pending transforms

[atbenmurray]

Describe the bug
Lazy resampling will execute the following pipeline incorrectly:

d = torch.zeros((1, 40, 40, 40))

xform = mt.Compose(
    [
        mt.Flip(spatial_axis=0),
        mt.Rotate90(k=1),
        mt.Rand2DElastic(3, (5, 10)),
        mt.Zoom(zoom=0.8),
        mt.Spacing(pixdim=1.2),
    ],
    lazy_evaluation=True,
    verbose=True
)
r = xform(d)

My understanding is that this was implemented deliberately to provide improved performance on some of the existing pipelines that are being benchmarked, but it will break a lot of our users pipelines if the non-lazy transforms rely on the pending transforms to have been run:
. Noise will undergo spatial distortion that was not intended
. Non-lazy spatial transforms will output the wrong result

While we want the best performance on our pipelines, we shouldn't expect our users to have to modify their pipelines to execute correctly.

The most obvious solution is a new flag for Compose. We've discussed it as a 1.3 feature but now we need it for this release.

class Compose(Randomizable, InvertibleTransform):
    def __init__(
        self,
        transforms: Sequence[Callable] | Callable | None = None,
        map_items: bool = True,
        unpack_items: bool = False,
        lazy: bool | None = False,
        options: ??? | None = None,
        overrides: dict | None = None,
        logger_name: str | None = None,

Options can be thought of like compile flags for a c++ compiler.
We should pick a way of specifying them that is extensible enough for our purpose and easy for the user:

"optimize: lazy_to_front, another_option: 3" # option string
["optimize: lazy_to_front", another_option: 3"] # option string array
{"optimize": "lazy_to_front", another_option: 3"} #json-like

*** Flags for 1.2 ***
We only need one or two flags for 1.2
. "lazy_to_front", "lazy_to_back"

These should physically reorder the transforms for execution. It can reshuffle the transforms in a way similar to RandomOrder.

@myron This should cover the benchmark performance that you need, yes?
@wyli This seems implementable in #6257 fairly simply. Given that RandomOrder works fine during inverse AFAIK, it means we also don't need to lose invert

[wyli]

Thanks and agree that adding an option will allow for different use cases, one of the use cases is to achieve fast but less precise results compared with the non-lazy...

{"optimize": "lazy_to_front", another_option: 3"} #json-like

Looks highly extensible

[myron]

hi @atbenmurray I don't understand the issue well

. Noise will undergo spatial distortion that was not intended
. Non-lazy spatial transforms will output the wrong result

which Noise? there is no intensity Noise transform here. or do you mean something else?

--
Also

 "lazy_to_front", "lazy_to_back"
These should physically reorder the transforms for execution.

Will it move all lazy transforms to the front? (then it won't be the right sequence of transforms ,) . If I understand it correctly - to move ALL lazy trasforms to front/back, then it will be applicable only to some special cases (and not in general), since we can't simply re-arrange transforms.

A common transform list I use is following (with Indentities for current "lazy_evaluation" version)

• LoadImageD
• CropForeground based on image
• Indentity(image, label)
• SpacingD (resample)
• NormalizeIntensityd #not sure if I need Identity before
• SpatialPadd (image, label)
• Indentity(label)
• RandCropByLabelClassesd (based on label classes)
• RandomAffine
• RandomFlip
• Indentity(image)
• RandomSmooth
• RandomIntensityNoise

I don't think we can simply re-arrange lazy (spatial) transforms to the front.

---

I personally think we should do lazy transform explicit. Anyone who wants to use them must create a separate list of spatial transforms with Compose( spatial_transforms, lazy=True), and then use that new compose_transform as a step between main transforms

[atbenmurray]

@myron thanks for the detailed response! In general, people shouldn't have to use this compose option. The primary purpose for it is to work around the situation we have now for the upcoming release. As I understand it, we have benchmarked pipelines that rely on the dev branch behaviour.

The dev branch produces incorrect output in the general case of having lazy and non lazy transforms mixed together, because, unless you use Identity/Identityd (becoming ApplyPending/ApplyPendingd) to enforce ordering, all non-lazy transforms will be executed first.

This decision was taken to improve performance on a benchmark that had a Lambda transform in it (as far as I am aware from the other conversations), but it will break our user's pipelines in the general case. move_lazy_first can't work for the above pipeline but move_lazy_last might.

I'm doing some experiments with your pipeline so that I fully understand the implications

@wyli I guess I really need to get a complete list of the pipelines this behaviour was done for so we can be sure that this fix would sort those pipelines for benchmarking while allowing default behaviour to not break our users' pipelines

[wyli]

I think we just use the option as you suggested to allow for the flexibility in general, the user can choose what to use for their own pipelines. it'll be impossible to have an exhaustive list of use cases.

by the way I disagree with "The dev branch behaviour has a major issue in that unless you use Identity/Identityd (becoming ApplyPending/ApplyPendingd) to enforce ordering", could you rephrase to avoid misleading info?

[atbenmurray]

I personally think we should do lazy transform explicit. Anyone who wants to use them must create a separate list of spatial > transforms with Compose( spatial_transforms, lazy=True), and then use that new compose_transform as a step between main transforms

The goal has always been to allow it to work without breaking anything in the general case. That way it is accessible just as a switch that can be flipped at the policy level rather than having to rewrite pipelines. The user is always then free to get very fancy with how they organise their transforms to get the extra benefit, but move_lazy_last, for example, would have that effect.

[atbenmurray]

by the way I disagree with "The dev branch behaviour has a major issue in that unless you use Identity/Identityd (becoming ApplyPending/ApplyPendingd) to enforce ordering", could you rephrase to avoid misleading info?

I've modified it to say 'will produce incorrect output'.

[atbenmurray]

@wyli I really need to see the pipelines that the dev branch behaviour was intended for. Surely you can point me at them. The flag isn't a fix unless we know it fixes those specific cases.

[wyli]

sure, currently the main cases that I'm aware of are:

• https://github.com/Project-MONAI/research-contributions/blob/18bb1d859fa65ad2221d2a6b6771e82b8e81ce89/auto3dseg/algorithm_templates/segresnet/scripts/segmenter.py#L200
• https://github.com/Project-MONAI/research-contributions/blob/main/auto3dseg/algorithm_templates/swinunetr/configs/transforms_train.yaml
• https://github.com/Project-MONAI/research-contributions/blob/main/auto3dseg/algorithm_templates/dints/configs/transforms_train.yaml
• https://github.com/Project-MONAI/tutorials/blob/main/modules/lazy_resampling_benchmark.ipynb
• https://github.com/Project-MONAI/tutorials/blob/main/modules/lazy_resampling_compose.ipynb
• https://github.com/Project-MONAI/tutorials/blob/main/modules/lazy_resampling_functional.ipynb
• Update wholeBody_ct_segmentation model-zoo#397

[atbenmurray]

Great! So if we eyeball these and verify that moving all the lazy transforms last is the equivalent of dev branch behaviour, then this is the right modification for the PR and we should be able to avoid needing to re-benchmark

[wyli]

I've modified it to say 'will produce incorrect output'.

I think I understand your point, but my point is that, there are 200+ transforms in monai, even if it's in non-lazy mode, you can easily compose certain sequences of some of them that may not work as you think (especially if you don't read the documentation)... we can't say that the monai codebase is incorrect or has major issue, right?

[wyli]

also, we'll never achieve exactly the same result by design, like the padding modes, output data types (quantisations), and interpolation modes. In most use cases, there's a trade-off for speed/interpolation quality/memory footprint, in my opinion, we just need more documentation about the details and let the user decides what transform sequence to use.

[atbenmurray]

I've modified it to say 'will produce incorrect output'.

I think I understand your point, but my point is that, there are 200+ transforms in monai, even if it's in non-lazy mode, you can easily compose certain sequences of some of them that may not work as you think (especially if you don't read the documentation)... we can't say that the monai codebase is incorrect or has major issue, right?

That's why the default behaviour should be as has always been stated in the design and presentations and discussions, i.e. that the ordering of transform execution is unmodified, and that any changes to transform ordering should come from option flags.
That way, the majority of people can always just switch the lazy flag on and have the best chance of it working as expected.

Anyway, our plan is to add the option flag which should allow the benchmarked pipelines to run as discussed, and have the default behaviour be to not modify transform order. There won't be any need to document anything as an issue in docstrings or release notes, because it won't be an issue, just an option that allows you to move all the lazy transforms last if you want to do so.

draft of a docstring fragment for compose options:

option 'move_lazy_last': This option reorders the transforms in a given pipeline so that they are executed after all non-lazy transforms. This can be helpful for performance and image quality, as it results in fewer resampling operations. Note that this option should not be used with pipelines that require ordering to be preserved, such as if you have a mixed sequence of lazy and non-lazy spatial transforms that must be performed in the specified order. If your pipeline uses ApplyPending transforms, they effectively split the list into sub-lists for the purposes of reordering and reordering only happens within the sub-list

[atbenmurray]

Ok, so after trying a couple of different interpretations of the dev lazy mode, I know what move_lazy_last means:

So:

'LX' means lazy transform with id X (id doesn't mean anything here, just helps with clarity in the examples)
'NY' means non-lazy transform with id Y
'AZ' means ApplyPending/ApplyPendingD with id Z

Given a list of transforms:

['LA', 'NB', 'LC', 'AD', 'LE', 'NF']

The presence of an ApplyPending splits the transform list into multiple sublists

['LA', 'NB', 'LC'], ['AD'], ['LE', 'NF']

Each sublist is then sorted so non-lazy transforms go first, lazy transforms go next, and then the ApplyPending goes last:

['NB', 'LA', 'LC'], ['AD'], ['NF', 'LE']

Which then concatenates back to:

['NB', 'LA', 'LC', 'AD', 'NF', 'LE']

This precisely mimics the behaviour in dev and makes sense from a user perspective; it is easy to explain that an ApplyPending effectively splits the list of transforms into sublists and schedules each separately.

I've implemented this and I am running a wide selection of tests on it.