fusing spatial transforms

[wyli]

Lazy resampling (fusing spatial transforms)

Is your feature request related to a problem? Please describe. Follow up of https://github.com/Project-MONAI/MONAI/discussions/4198, https://github.com/Project-MONAI/MONAI/issues/112, when there are various spatial transforms such as Spacing RandZoom being used in Compose, automatically fusing them to run one step resampling would speed up the process and reduce the resampling errors by repeated resamplings.

Introduction

Lazy resampling (or fusing spatial transforms) is a feature that aims to eliminate unnecessary resample steps from preprocessing pipelines containing multiple spatial transforms. It does so by composing those steps in a manner that reflects established best practice from the world of graphics libraries.

The goal is that this functionality is incorporated into the implementation of Compose and the spatial transforms themselves, so that a user can get the benefit of this feature in a way that is completely compatible code-wise with previous versions of MONAI.

Historically speaking, this functionality has existed only within the scope of Affine and Affined, which enable multiple affine transforms in a predetermined order within a single transform, in #112.

Benefits

Lazy resampling has the following benefits:

• It improves the efficiency of a preprocessing pipeline in terms of memory and time, particularly the latter
• It improves the quality of the resulting images, reducing issues with padding
• It enables reordering of pipelines to simplify reasoning about pipelines#5249
• It enables simpler implementations of transforms

Minor refactor vs. major refactor

There are two ongoing PRs relating to this feature:

• #4911 is a minimal modification to the existing code-base that seeks to enable the feature without a major modification to the way that transforms are implemented
• #4922 is a major refactor that seeks to re-implement spatial transforms in a simplified manner, along with a more extensively refactored Compose to consolidate areas of the code that manipulate transform pipelines (such as CachedDataset)

Impact on codebase

Transforms

All spatial transforms must be modified to be resampling aware. This can be done in a minimal way (#4911) or in a more extensive way (#4922) that decouples the transform itself from the resulting resample operation.

Compose

Compose must be modified to handle lazy resampling although thee are a number of ways this can be achieved

Resample (optional)

Resample can optionally be modified in order to reduce the complexity of spatial transforms, according to #5010. This allows all spatial transforms to be implemented in terms of a description of the transform while still being able to take advantage of lower-cost resampling techniques when possible.

Complications

Performing lazy resampling across multiple transforms adds some feature complexity that must be understood and overcome:

• Modifications to Compose
• Parameter compatibility
• Other code that modifies the transform list
  • Dataset caching
  • Processing the inverse
  • Histogram-sampling transforms
• Multi-sample transforms
• Compiled transforms

Modifications to Compose

The goal with this functionality as that the user should not need to alter existing code in order for lazy resampling to work; i.e. that it should be the default behaviour moving forward. As such, additional work must be done behind the scenes in order for a pipeline to be lazily executed. This can manifest in code in a number of different ways:

• Modification to Compose: Compose can have the concept of lazy evaluation built into it, so that if it determines that there are one or more transforms capable of lazy execution, that it causes resampling to occur in the appropriate locations
• Introduction of Pipeline Compilation: A pipeline compiler that takes a list of transforms and generates an augmented / modified list of transforms that can be executed in its place. This approach might consolidate other places in the codebase that execute a modified list of transforms (see below).

Parameter compatibility

There are multiple parameters that can be set across most spatial transforms, such as mode, padding_mode and dtype. If these are incompatible between transforms, then resampling must happen immediately rather than deferring it to later in the pipeline.

Other code that modifies the transform list

There are a number of features that cause a list of transforms to be modified before / during their execution:

• Dataset caching: CachedDataset and PersistentDataset both implement a scheme whereby deterministic transforms that happen before the initial randomized transform are cached to memory / drive, and skipped subsequently
• Inverse: Inverting the transforms runs the transforms backwards using the invert method rather than __call__

With the need for lazy resampling to also execute a modified list of transforms, there is significant potential for clashes if these list altering mechanisms aren't designed to work with each other.

Histogram-sampling transforms

Some transforms perform patch-based sampling given histograms of data; typically labels that accompany image or volume data. This might require an immediate resample for if the patch based transform follows other spatial transforms, but this resample may then subsequently be thrown out in favour of a single true resample step at the end of the spatial transforms.

Multi-sample transforms

Some transforms allow for multiple samples to be performed, such as when multiple random patches are selected from a data sample. Lazy resampling causes complications relating to MetaTensor, as there isn't strictly speaking a need for a new metatensor instance until resampling occurs, but MetaTensor is not designed to hold multiple sets of pending transforms. There are several ways that this can be handled:

• Run each sample through the pipeline to the point where resampling occurs in a depth first manner
• Modify MetaTensor to handle arrays of pending transforms
• Allow MetaTensor instances to share (in a read only fashion) the underlying tensor data between multiple meta tensor instances. This is currently the preferred approach

Compiled transforms

If pipeline compilation is chosen, it will be necessary to write new transforms that implement the modified behaviour. These include:

• Transforms that check for cached data and execute if that cached data is not present. These will typically contain a set of transforms from the original pipeline.
• Apply / Applyd required to execute lazy resampling

[atbenmurray]

@wyli Do you mind if I add some more detail to this feature request?

[wyli]

no, please feel free to add more info @atbenmurray

[myron]

This is great effort, to interpolate only once, it would avoid the repeated interpolation artifacts and should speed up the transforms pipeline, and use less memory, in theory.

At the same time, it's very difficult to implement for any general transforms, and may result in many potential bugs. Perhaps we can start by supporting only a subset of transforms (Spacing, Cropping and fused Affine (which includes Rotation, Scaling, Affine)).

And may be it's better for a user to explicitly indicate which transforms to fuse together, e.g.

fused_transform = transforms.FusedSpatialTranform([SpacingD(), RandomCroppingD(), RandAffineD()])

transforms = transforms.Compose[LoadImageD(), ToTensor(), fused_transform,  NormalizeIntensities()]

It's one extra step for the user, but it could be easier to test, and returns the error right away if the transforms can not be fused.

In this case each transform can have a method to accept a grid (X, Y, Z coordinates) for interpolation, and each transform will update it (e.g. crop grid, or rotate grid or scale).

and the transforms.FusedSpatialTranform() class will call all those spatial grid manipulations and interpolate once at the end

[atbenmurray]

This is great effort, to interpolate only once, it would avoid the repeated interpolation artifacts and should speed up the transforms pipeline, and use less memory, in theory.

At the same time, it's very difficult to implement for any general transforms, and may result in many potential bugs. Perhaps we can start by supporting only a subset of transforms (Spacing, Cropping and fused Affine (which includes Rotation, Scaling, Affine)).

And may be it's better for a user to explicitly indicate which transforms to fuse together, e.g.

fused_transform = transforms.FusedSpatialTranform([SpacingD(), RandomCroppingD(), RandAffineD()])

transforms = transforms.Compose[LoadImageD(), ToTensor(), fused_transform,  NormalizeIntensities()]

It's one extra step for the user, but it could be easier to test, and returns the error right away if the transforms can not be fused.

In this case each transform can have a method to accept a grid (X, Y, Z coordinates) for interpolation, and each transform will update it (e.g. crop grid, or rotate grid or scale).

and the transforms.FusedSpatialTranform() class will call all those spatial grid manipulations and interpolate once at the end

The goals has been that lazy resampling works without you having to change the code at all. As such, transforms that are able to execute lazily will do so unless the user turns it off. It isn't an issue for spatial transforms in general; any affine transformation can be composed with others to create a cumulative transform, and crops / pads can be rephrased as translation operations that change the region of interest.

Transforms expressed as grids (elastic, etc.) can also be composed with affine transforms. It is only when we need to compose a grid with a grid that we need to pause and resample before continuing.

When transforms can't be fused, we just perform a resample and then move on to the rest of the transforms, so only transforms that are explicitly lazy in combination with each other will result in intermediate resamples.

Does that clarify how things work in lazy resampling and / or ameliorate your concerns?

[myron]

@atbenmurray thank you for the reply. I understand the idea of lazy resampling and fusing transforms. And I understand the idea of hiding this from the user, to simplify things for the user.

But to me it seems better to explicitly define which transforms to fuse, and separate this logic, for the following reasons

• If the user is doing this transform sequence: [SpatialResampling->RandomGaussianSmoothing->RandomAffine]. Then we will not be able to fuse anything. There will be no performance gains as you need to resample 2 times. And a user would not know it. We could put in the documentation to try to stack a certain transforms together first, but not too many people will figure it out.
  Instead if a user explicitly calls transform.FusedSpatialTranform([SpatialResampling(), RandomGaussianSmoothing(), RandomAffine()]), we can show an error, that they can not be fused, forcing a use to rearrange transform.FusedSpatialTranform([SpatialResampling(), RandomAffine()]), followed by RandomGaussianSmoothing()
• Another advantage to do it this way would be for a class FusedSpatialTranform to maintain the extra info (a sequence of transforms, grid changed to perform) and to perform the final interplation. And in this class we can define which interpolation type to use, and which boundary condition, extrapolation values to use. this will be easier to debug too.

these are my thoughts, but there are multiple equally good ways/solutions to accomplish this (with it's own challenges).