Merge branch 'dev' into warp-doc-differentiability-warning

Author: ericspod

monai/apps/pathology/transforms/stain/array.py

@@ -85,7 +85,12 @@ def _deconvolution_extract_stain(self, image: np.ndarray) -> np.ndarray:
         v_max = eigvecs[:, 1:3].dot(np.array([(np.cos(max_phi), np.sin(max_phi))], dtype=np.float32).T)
 
         # a heuristic to make the vector corresponding to hematoxylin first and the one corresponding to eosin second
-        if v_min[0] > v_max[0]:
+        # Hematoxylin: high blue, lower red (low R/B ratio)
+        # Eosin: high red, lower blue (high R/B ratio)
+        eps = np.finfo(np.float32).eps
+        v_min_rb_ratio = v_min[0] / (v_min[2] + eps)
+        v_max_rb_ratio = v_max[0] / (v_max[2] + eps)
+        if v_min_rb_ratio < v_max_rb_ratio:
             he = np.array((v_min[:, 0], v_max[:, 0]), dtype=np.float32).T
         else:
             he = np.array((v_max[:, 0], v_min[:, 0]), dtype=np.float32).T

monai/networks/blocks/activation_checkpointing.py

@@ -0,0 +1,41 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+
+class ActivationCheckpointWrapper(nn.Module):
+    """Wrapper applying activation checkpointing to a module during training.
+
+    Args:
+        module: The module to wrap with activation checkpointing.
+    """
+
+    def __init__(self, module: nn.Module) -> None:
+        super().__init__()
+        self.module = module
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with optional activation checkpointing.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            Output tensor from the wrapped module.
+        """
+        return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))

monai/networks/nets/unet.py

@@ -17,11 +17,12 @@
 import torch
 import torch.nn as nn
 
+from monai.networks.blocks.activation_checkpointing import ActivationCheckpointWrapper
 from monai.networks.blocks.convolutions import Convolution, ResidualUnit
 from monai.networks.layers.factories import Act, Norm
 from monai.networks.layers.simplelayers import SkipConnection
 
-__all__ = ["UNet", "Unet"]
+__all__ = ["UNet", "Unet", "CheckpointUNet"]
 
 
 class UNet(nn.Module):
@@ -298,4 +299,29 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+class CheckpointUNet(UNet):
+    """UNet variant that wraps internal connection blocks with activation checkpointing.
+
+    See `UNet` for constructor arguments. During training with gradients enabled,
+    intermediate activations inside encoder-decoder connections are recomputed in
+    the backward pass to reduce peak memory usage at the cost of extra compute.
+    """
+
+    def _get_connection_block(self, down_path: nn.Module, up_path: nn.Module, subblock: nn.Module) -> nn.Module:
+        """Returns connection block with activation checkpointing applied to all components.
+
+        Args:
+            down_path: encoding half of the layer (will be wrapped with checkpointing).
+            up_path: decoding half of the layer (will be wrapped with checkpointing).
+            subblock: block defining the next layer (will be wrapped with checkpointing).
+
+        Returns:
+            Connection block with all components wrapped for activation checkpointing.
+        """
+        subblock = ActivationCheckpointWrapper(subblock)
+        down_path = ActivationCheckpointWrapper(down_path)
+        up_path = ActivationCheckpointWrapper(up_path)
+        return super()._get_connection_block(down_path, up_path, subblock)
+
+
 Unet = UNet

tests/apps/pathology/transforms/test_pathology_he_stain.py

@@ -48,7 +48,7 @@
 # input pixels not uniformly filled, leading to two different stains extracted
 EXTRACT_STAINS_TEST_CASE_5 = [
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[0.70710677, 0.18696113], [0.0, 0.0], [0.70710677, 0.98236734]]),
+    np.array([[0.18696113, 0.70710677], [0.0, 0.0], [0.98236734, 0.70710677]]),
 ]
 
 # input pixels all transparent and below the beta absorbance threshold
@@ -68,7 +68,7 @@
 NORMALIZE_STAINS_TEST_CASE_4 = [
     {"target_he": np.full((3, 2), 1)},
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]]]),
+    np.array([[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]]]),
 ]
 
 
@@ -135,7 +135,7 @@ def test_result_value(self, image, expected_data):
           [[0.18696113],[0],[0.98236734]] and
           [[0.70710677],[0],[0.70710677]] respectively
         - the resulting extracted stain should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]]
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]]
         """
         if image is None:
             with self.assertRaises(TypeError):
@@ -206,17 +206,17 @@ def test_result_value(self, arguments, image, expected_data):
 
         For test case 4:
         - For this non-uniformly filled image, the stain extracted should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]], as validated for the
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]], as validated for the
           ExtractHEStains class. Solving the linear least squares problem (since
           absorbance matrix = stain matrix * concentration matrix), we obtain the concentration
-          matrix that should be [[-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508],
-          [5.8022, 0, 0, 0, 0, 0]]
+          matrix that should be [[5.8022, 0, 0, 0, 0, 0],
+          [-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508]]
         - Normalizing the concentration matrix, taking the matrix product of the
           target stain matrix and the concentration matrix, using the inverse
           Beer-Lambert transform to obtain the RGB image from the absorbance
           image, and finally converting to uint8, we get that the stain normalized
-          image should be [[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]],
-          [[33, 33, 33], [33, 33, 33]]]
+          image should be [[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]],
+          [[85, 85, 85], [85, 85, 85]]]
         """
         if image is None:
             with self.assertRaises(TypeError):

tests/apps/pathology/transforms/test_pathology_he_stain_dict.py

@@ -42,7 +42,7 @@
 # input pixels not uniformly filled, leading to two different stains extracted
 EXTRACT_STAINS_TEST_CASE_5 = [
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[0.70710677, 0.18696113], [0.0, 0.0], [0.70710677, 0.98236734]]),
+    np.array([[0.18696113, 0.70710677], [0.0, 0.0], [0.98236734, 0.70710677]]),
 ]
 
 # input pixels all transparent and below the beta absorbance threshold
@@ -62,7 +62,7 @@
 NORMALIZE_STAINS_TEST_CASE_4 = [
     {"target_he": np.full((3, 2), 1)},
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]]]),
+    np.array([[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]]]),
 ]
 
 
@@ -129,7 +129,7 @@ def test_result_value(self, image, expected_data):
           [[0.18696113],[0],[0.98236734]] and
           [[0.70710677],[0],[0.70710677]] respectively
         - the resulting extracted stain should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]]
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]]
         """
         key = "image"
         if image is None:
@@ -200,17 +200,17 @@ def test_result_value(self, arguments, image, expected_data):
 
         For test case 4:
         - For this non-uniformly filled image, the stain extracted should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]], as validated for the
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]], as validated for the
           ExtractHEStains class. Solving the linear least squares problem (since
           absorbance matrix = stain matrix * concentration matrix), we obtain the concentration
-          matrix that should be [[-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508],
-          [5.8022, 0, 0, 0, 0, 0]]
+          matrix that should be  [[5.8022, 0, 0, 0, 0, 0],
+          [-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508]]
         - Normalizing the concentration matrix, taking the matrix product of the
           target stain matrix and the concentration matrix, using the inverse
           Beer-Lambert transform to obtain the RGB image from the absorbance
           image, and finally converting to uint8, we get that the stain normalized
-          image should be [[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]],
-          [[33, 33, 33], [33, 33, 33]]]
+          image should be [[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]],
+          [[85, 85, 85], [85, 85, 85]]]
         """
         key = "image"
         if image is None:

tests/networks/nets/test_checkpointunet.py

@@ -0,0 +1,186 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks import eval_mode
+from monai.networks.layers import Act, Norm
+from monai.networks.nets.unet import CheckpointUNet, UNet
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+TEST_CASE_0 = [  # single channel 2D, batch 16, no residual
+    {
+        "spatial_dims": 2,
+        "in_channels": 1,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 0,
+    },
+    (16, 1, 32, 32),
+    (16, 3, 32, 32),
+]
+
+TEST_CASE_1 = [  # single channel 2D, batch 16
+    {
+        "spatial_dims": 2,
+        "in_channels": 1,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 1,
+    },
+    (16, 1, 32, 32),
+    (16, 3, 32, 32),
+]
+
+TEST_CASE_2 = [  # single channel 3D, batch 16
+    {
+        "spatial_dims": 3,
+        "in_channels": 1,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 1,
+    },
+    (16, 1, 32, 24, 48),
+    (16, 3, 32, 24, 48),
+]
+
+TEST_CASE_3 = [  # 4-channel 3D, batch 16
+    {
+        "spatial_dims": 3,
+        "in_channels": 4,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 1,
+    },
+    (16, 4, 32, 64, 48),
+    (16, 3, 32, 64, 48),
+]
+
+TEST_CASE_4 = [  # 4-channel 3D, batch 16, batch normalization
+    {
+        "spatial_dims": 3,
+        "in_channels": 4,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 1,
+        "norm": Norm.BATCH,
+    },
+    (16, 4, 32, 64, 48),
+    (16, 3, 32, 64, 48),
+]
+
+TEST_CASE_5 = [  # 4-channel 3D, batch 16, LeakyReLU activation
+    {
+        "spatial_dims": 3,
+        "in_channels": 4,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 1,
+        "act": (Act.LEAKYRELU, {"negative_slope": 0.2}),
+        "adn_ordering": "NA",
+    },
+    (16, 4, 32, 64, 48),
+    (16, 3, 32, 64, 48),
+]
+
+TEST_CASE_6 = [  # 4-channel 3D, batch 16, LeakyReLU activation explicit
+    {
+        "spatial_dims": 3,
+        "in_channels": 4,
+        "out_channels": 3,
+        "channels": (16, 32, 64),
+        "strides": (2, 2),
+        "num_res_units": 1,
+        "act": (torch.nn.LeakyReLU, {"negative_slope": 0.2}),
+    },
+    (16, 4, 32, 64, 48),
+    (16, 3, 32, 64, 48),
+]
+
+CASES = [TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6]
+
+
+class TestCheckpointUNet(unittest.TestCase):
+    @parameterized.expand(CASES)
+    def test_shape(self, input_param, input_shape, expected_shape):
+        """Validate CheckpointUNet output shapes across configurations.
+
+        Args:
+            input_param: Dictionary of UNet constructor arguments.
+            input_shape: Tuple specifying input tensor dimensions.
+            expected_shape: Tuple specifying expected output tensor dimensions.
+        """
+        net = CheckpointUNet(**input_param).to(device)
+        with eval_mode(net):
+            result = net.forward(torch.randn(input_shape).to(device))
+            self.assertEqual(result.shape, expected_shape)
+
+    def test_checkpointing_equivalence_eval(self):
+        """Confirm eval parity when checkpointing is inactive."""
+        params = dict(
+            spatial_dims=2, in_channels=1, out_channels=2, channels=(8, 16, 32), strides=(2, 2), num_res_units=1
+        )
+
+        x = torch.randn(2, 1, 32, 32, device=device)
+
+        torch.manual_seed(42)
+        net_plain = UNet(**params).to(device)
+
+        torch.manual_seed(42)
+        net_ckpt = CheckpointUNet(**params).to(device)
+
+        # Both in eval mode disables checkpointing logic
+        with eval_mode(net_ckpt), eval_mode(net_plain):
+            y_ckpt = net_ckpt(x)
+            y_plain = net_plain(x)
+
+        # Check shape equality
+        self.assertEqual(y_ckpt.shape, y_plain.shape)
+
+        # Check numerical equivalence
+        self.assertTrue(
+            torch.allclose(y_ckpt, y_plain, atol=1e-6, rtol=1e-5),
+            f"Eval-mode outputs differ: max abs diff={torch.max(torch.abs(y_ckpt - y_plain)).item():.2e}",
+        )
+
+    def test_checkpointing_activates_training(self):
+        """Verify checkpointing recomputes activations during training."""
+        params = dict(
+            spatial_dims=2, in_channels=1, out_channels=1, channels=(8, 16, 32), strides=(2, 2), num_res_units=1
+        )
+
+        net = CheckpointUNet(**params).to(device)
+        net.train()
+
+        x = torch.randn(2, 1, 32, 32, device=device, requires_grad=True)
+        y = net(x)
+        loss = y.mean()
+        loss.backward()
+
+        # gradient flow check
+        grad_norm = sum(p.grad.abs().sum() for p in net.parameters() if p.grad is not None)
+        self.assertGreater(grad_norm.item(), 0.0)
+
+
+if __name__ == "__main__":
+    unittest.main()