Lazy eval refactor

docs/source/lazy_resampling.rst

@@ -0,0 +1,104 @@
+:github_url: https://github.com/Project-MONAI/MONAI
+
+Lazy Resampling
+===============
+
+.. toctree::
+   : maxdepth: 2
+
+    mb_specification
+    config_syntax.md
+
+Introduction
+^^^^^^^^^^^^
+
+Lazy Resampling is a new feature for MONAI 1.2. This feature is still experimental at this time and it is possible that
+behaviour and APIs will change in upcoming releases.
+
+Lazy resampling is a feature that can be used to improve preprocessing pipelines in the following ways:
+ * it can improve pipeline execution time
+ * it can improve pipeline memory usage
+ * it can improve image and segmentation quality by reducing incidental noise caused by resampling
+
+How Lazy Resampling changes preprocessing
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In order to understand how lazy resampling changes preprocessing, we'll first discuss standard processing pipeline
+behaviour, and then compare it with the way lazy resampling works.
+
+Traditional resampling pipelines
+++++++++++++++++++++++++++++++++
+
+With traditional resampling, found both in MONAI and many other preprocessing libraries, you typically define a sequence
+of transforms and pass them to a ``Compose`` object, such as `monai.transforms.compose.Compose`_.
+
+Example::
+
+    transforms = [
+        LoadImaged(keys=["img", "seg"], ...),
+        EnsureChannelFirstd(keys=["img", "seg"], ...),
+        Spacingd(keys=["img", "seg"], ...),
+        Orientationd(keys=["img", "seg"], ...),
+        RandSpatialCropd(keys=["img", "seg"], ...),
+        RandRotate90d(keys=["img", "seg"], ...),
+        RandRotated(keys=["img", "seg"], ...),
+        RandZoomd(keys=["img", "seg"], ...),
+        RandGaussianNoised(keys="img", ...),
+    ]
+    compose = Compose(transforms)
+
+    # elsewhere this will be called many times (such as in a Dataset instance)
+    outputs = compose(inputs)
+::
+
+The following will then happen when we call ``compose(inputs)``:
+
+1. ``LoadImaged`` is called with its inputs (a dictionary of strings containing file locations). This loads and
+   returns a dictionary of the corresponding data samples
+2. ``EnsureChannelFirstd`` is called with the dictionary of data samples and adds a channel so that they have the
+   appropriate shape for the rest of the pipeline
+3. ``Spacingd`` is called and reinterpolates the data samples
+4. ``Orientationd`` permutes the data samples so that their spatial dimensions are reorganised
+5. ``RandSpatialCropd`` crops a random patch of the data samples, throwing away the rest of the data in the process
+6. ``RandRotate90d`` has a chance of performing a tensor-based rotation of the data samples
+7. ``RandRotated`` has a chance of performing a full resample of the data samples
+8. ``RandZoomd`` has a chance of performing a reinterpolation of the data samples
+9. ``RandGaussianNoised`` has a chance of adding noise to ``img``
+
+Overall, there are up to three occasions where the data is either interpolated or resampled through spatial transforms.
+Furthermore, the crop that occurs means that the output data samples might contain pixels for which there is data but
+that show padding values, because the data was thrown away by ``RandSpatialCrop``.
+
+Each of these operations takes time and memory, but, as we can see in the example above, also creates resampling
+artifacts and can even destroy data in the resulting data samples (see `lazy resampling best practices`_ for examples).
+
+Lazy resampling pipelines
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Lazy resampling works very differently. When you execute the same pipeline with `lazy=True`, the following happens:
+
+1. ``LoadImaged`` behaves identically
+2. ``EnsureChannelFirstd`` behaves identically
+3. ``Spacingd`` is executing lazily. It puts a description of the operation that it wants to perform onto a list of
+   pending operations
+4. ``Orientationd`` is executing lazily. It adds a description of its own operation to the pending operation list so
+   now there are 2 pending operations
+5. ``RandSpatialCropd`` is executing lazily. It adds a description of its own operation to the pending operation list
+   so now there are 3 pending operations
+6. ``RandRotate90d`` is executing lazily. It adds a description of its own operation to the pending operation list
+   so now there are 4 pending operations
+7. ``RandRotated`` is executing lazily. It adds a description of its own operation to the pending operation list
+   so now there are 5 pending operations
+8. ``RandZoomd`` is executing lazily. It adds a description of its own operation to the pending operation list
+   so now there are 6 pending operations
+   1. ``[Spacingd, Orientationd, RandSpatialCropd, RandRotate90d, RandRotated, RandZoomd]`` are all on the pending
+      operations list but have yet to be carried out on the data
+9. ``RandGaussianNoised`` is not a lazy transform. It is now time for the pending operations to be evaluated. Their
+   descriptions are mathematically composited together, to determine the operation that results from all of them
+   being carried out. This is then applied in a single resample operation. Once that is done, ``RandGaussianNoised``
+   operates on the resulting data
+
+The single resampling operation has less noise induced by resampling, as it only occurs once in this pipeline rather
+than three times in the traditional pipeline. More importantly, although the crop describes an operation to keep only a
+subset of the data sample, the crop is not performed until after the spatial transforms are completed, which means that
+all of the data sample that is within bounds is preserved and is part of the resulting output.

docs/source/transforms.rst

@@ -958,6 +958,17 @@ MRI Transforms
     :special-members: __call__
 
 
+Lazy
+^^^^
+
+`ApplyPending`
+""""""""""""""
+
+.. autoclass:: ApplyPending
+    :members:
+    :special-members: __call__
+
+
 Utility
 ^^^^^^^
 
@@ -1912,6 +1923,17 @@ Smooth Field (Dict)
     :special-members: __call__
 
 
+Lazy (Dict)
+^^^^^^^^^^^
+
+`ApplyPendingd`
+"""""""""""""""
+
+.. autoclass:: ApplyPendingd
+    :members:
+    :special-members: __call__
+
+
 Utility (Dict)
 ^^^^^^^^^^^^^^
 
@@ -2211,9 +2233,3 @@ Utilities
 
 .. automodule:: monai.transforms.utils_pytorch_numpy_unification
     :members:
-
-Lazy
-----
-.. automodule:: monai.transforms.lazy
-    :members:
-    :imported-members:

monai/apps/auto3dseg/data_analyzer.py

@@ -212,7 +212,7 @@ def get_all_case_stats(self, key="training", transform_list=None):
                 manager_list = manager.list()
                 processes = []
                 for rank in range(nprocs):
-                    p = tmp_ctx.Process(
+                    p = tmp_ctx.Process(  # type: ignore[attr-defined]
                         target=self._get_all_case_stats, args=(rank, nprocs, manager_list, key, transform_list)
                     )
                     processes.append(p)

monai/apps/detection/transforms/dictionary.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import warnings
 from collections.abc import Hashable, Mapping, Sequence
 from copy import deepcopy
 from typing import Any
@@ -1326,7 +1327,11 @@ def __init__(
         super().__init__(self.image_keys + self.box_keys, prob, max_k, spatial_axes, allow_missing_keys)
         self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys))
 
-    def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Mapping[Hashable, torch.Tensor]:
+    def __call__(
+        self, data: Mapping[Hashable, torch.Tensor], lazy: bool | None = None
+    ) -> Mapping[Hashable, torch.Tensor]:
+        if lazy is True:
+            warnings.warn("RandRotateBox90d cannot currently execute lazily; ignoring lazy=True")
         self.randomize()
         d = dict(data)
 

monai/apps/reconstruction/transforms/dictionary.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import warnings
 from collections.abc import Hashable, Mapping, Sequence
 
 import numpy as np
@@ -213,10 +214,10 @@ class ReferenceBasedSpatialCropd(Cropd):
     """
 
     def __init__(self, keys: KeysCollection, ref_key: str, allow_missing_keys: bool = False) -> None:
-        super().__init__(keys, cropper=None, allow_missing_keys=allow_missing_keys)  # type: ignore
+        super().__init__(keys, cropper=None, allow_missing_keys=allow_missing_keys, lazy=False)  # type: ignore
         self.ref_key = ref_key
 
-    def __call__(self, data: Mapping[Hashable, Tensor]) -> dict[Hashable, Tensor]:
+    def __call__(self, data: Mapping[Hashable, Tensor], lazy: bool | None = None) -> dict[Hashable, Tensor]:
         """
         This transform can support to crop ND spatial (channel-first) data.
         It also supports pseudo ND spatial data (e.g., (C,H,W) is a pseudo-3D
@@ -229,6 +230,9 @@ def __call__(self, data: Mapping[Hashable, Tensor]) -> dict[Hashable, Tensor]:
         Returns:
             the new data dictionary
         """
+        if lazy is True:
+            warnings.warn("ReferenceBasedSpatialCropd cannot currently execute lazily; ignoring lazy=True")
+
         d = dict(data)
 
         # compute roi_size according to self.ref_key

monai/transforms/__init__.py

@@ -230,7 +230,9 @@
 from .inverse_batch_transform import BatchInverseTransform, Decollated, DecollateD, DecollateDict
 from .io.array import SUPPORTED_READERS, LoadImage, SaveImage
 from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict
-from .lazy.functional import apply_transforms
+from .lazy.array import ApplyPending
+from .lazy.dictionary import ApplyPendingd, ApplyPendingD, ApplyPendingDict
+from .lazy.functional import apply_pending
 from .lazy.utils import combine_transforms, resample
 from .meta_utility.dictionary import (
     FromMetaTensord,