Lazy Resampling Roadmap (1.4+)

[atbenmurray]

This feature request details the proposed extensions to the Lazy Resampling feature introduced in 1.2. It is anticipated that some subset of these proposed extensions will be implemented over the next several releases of MONAI, and should be done so in a modular fashion.

Consolidated Resampler

The current implementation of the consolidated resample is still a partial solution that is missing both the ability to map between analogous modes and padding modes and the ability to always run the best kind of transform. There is a partial implementation of this but it is limited to tensor ops or general grid based resampling; it should be possible to make use of interpolation also. This additionally provides the functionality to determine whether a resample is required due to a particular mode not being available to a particular resample operation.

Consolidated Resampler - non-grid affine resample
All general resampling is currently done with grids. This imposes both a processing and memory overhead for preprocessing, as the majority of spatial transforms do not require grids. CPU and GPU implementations of resampling without grids completes the work for consolidated resampling

Pure lazy spatial transforms
Once a basic Consolidated Resampler is available, it becomes easier to move to pure lazy spatial transforms. Pure lazy spatial transforms provide a simpler way to specify transforms, as the transform is only concerned with the definition of the operation, not its execution. In non-lazy operation, the defined operation is immediately carried out. In lazy mode, it is allowed to accumulate. This reduces the code footprint and complexity of our current transform implementations.

Pure lazy cropping transforms
Similar to spatial transforms, crops can be defined solely in terms of affine matrices plus metadata.

Lazy inverse
Lazy inverse provides a default mechanism for inverting lazy transforms that can be defined in LazyTransform. Pure, invertible lazy transforms can all make use of this same mechanism, all but eliminating custom invert code, excepting for transforms that must perform a custom invert operation. The key point is that, if lazy transforms are pure (i.e. they define their transform completely through affine transformation + metadata) then all inversion is done by simply inverting the transform.

Advanced lazy inverse
Transforms are accumulated during a forward pass but the resulting composed transforms are thrown away. These could be cached and used to further optimise the inverse pass if required. It is not immediately obvious that this would provide a significant performance boost, but can be further investigated if scenarios arise where this would be beneficial.

Pipeline Compiler and structural transforms
The original design for Lazy Resampling used a pipeline compiler to take a set of transforms and reorganise them for lazy resampling. The resulting pipeline doesn't need any special handling by a Compose object. ApplyPending transforms are part of this design, the idea being that the pipeline compiler inserts them where they are needed. This idea expands to other concepts where one transforms a pipeline once rather than evaluating the pipeline in a potentially complex fashion every time it is executed. Structural Transforms are a way of decoupling and conquering this complexity. Examples include:
. CachingTransform: given a pipeline to be run before a caching mechanism and a pipeline to be run after a caching mechanism, check whether a key is present in the cache. If it isn't, run the pre caching pipeline and cache it, otherwise, fetch the output from the cache. Then run the post caching pipeline
. DepthFirstMultiSampler: given a transform that produces multiple samples and a pipeline downstream of that transform, iterate over the output of the multii-sample transform and execute the rest of the pipeline in a depth-first fashion

This has a number of benefits:
. The pipeline compiler can be used internally by MONAI's Compose implementation and explicitly by users who want to do lazy resampling and externally by users who want (or need) to use a more generic Compose implementation
. The pipeline compiler allows us to avoid implementing lazy resampling (and other complexity adding features) explicitly in the compose pipeline and having to perform the same checks every time the pipeline is executed
. Depth first multi-sampling (see below)
. The pipeline compiler can also be used to cleanly implement caching dataset mechanisms

Depth first multi-sampling
Multi-sampling is currently performed breadth-first. This means that, when a multi-sampling transform is encountered, it generates all of its samples and returns them in a list. In traditional resampling, this can dramatically increase memory usage. Lazy resampling mitigates some of the memory increase as the resulting metatensors share the underlying data, but this is still suboptimal from a memory standpoint if lazy transforms generate a grid rather than a matrix, and we still hit a memory spike at the point the downstream samples undergo a resample operation. Moving to a generator style multi-sampling mechanism means that, regardless of whether transforms are lazy, the cost of their transform representation, or the number of samples to be generated, memory can remain constant over multi-sampling. This is very easy to implement if we have the Pipeline Compiler

Full viewport calculation
At the moment, viewports are not properly calculated and maintained through lazy composition, instead, the resulting image shape is transformed by the calculations and passed forward, but it is always rounded to integer values. It would be better to maintain a full viewport of floating point values and then determine the resulting shape from this at the point the transform is being applied.

Lazy execution for transforms that access data
There are a number of cropping transforms that look at data values in order to determine what crop to make. As it stands, these force execution of incoming data, and can only be lazy with their outputs. We can modify lazy execution such that, although the forced execution happens, we can throw away that data after it is analysed and continue the lazy execution for the output of the pipeline. This mitigates data loss from crop-type resamples. Because of its effect on performance, we might want to make this an optional flag for lazy execution.

Lazy noise
At the moment, noise transforms are not able to be applied lazily. We have in the past considered moving the noise transforms to after (or before) the spatial transforms but that is also suboptimal as the spatial transforms distort the noise once applied that moving the noise transforms breaks that. However, we can still make noise lazy by specifying noise as a pending operation. When noise is specified in the pipeline, we can track the spatial transforms that occur between the introduction of that noise and the point at which a resample is performed. We can sample the noise, then transform it with the necessary spatial transforms, and then apply it to the tensor being resampled.

Lazy matrix -> grid operations
Transforms that output grids (i.e. those that carry out deformations of various kinds) are not currently lazy, and should be made lazy. In this initial phase, they can accept inputs lazily but will resample on their output. This can be subsequently extended to grid -> matrix and grid -> grid operations.

Lazy grid -> matrix operations
Transforms that output grids can be positioned in an intermediate location in a sequence of lazy transforms. At composition time, each point in the grid can subsequently be transformed by the matrix, producing an updated grid as the result of the composition

Lazy grid -> grid operations
Transforms that output grids can also be lazily evaluated. At composition time, each point in the first grid can be vector multiplied with the point in the second grid, producing an updated grid as the result of the composition.

[wyli]

no action item, so I'm moving this to the discussion tab.

[vikashg]

@vikashg