Lazy Resampling Roadmap (1.4+) #7150
Description
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.
Each of these items should undergo a design discussion and be turned into separate feature requests once agreed upon.
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.
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.