fix: simplify test and make sure that checkpoint unet runs well in training

Fabio Ferreira · Fabio Ferreira · commit 1aa8e3c3ab45 · 2025-11-07T14:01:05.000Z
diff --git a/tests/networks/nets/test_checkpointunet.py b/tests/networks/nets/test_checkpointunet.py
@@ -1,6 +1,6 @@
 # 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 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
@@ -17,13 +17,12 @@
 from parameterized import parameterized
 
 from monai.networks import eval_mode
-from monai.networks.layers import Act, Norm
-from monai.networks.nets.unet import CheckpointUNet
+from monai.networks.nets.unet import CheckpointUNet, UNet
 from tests.test_utils import test_script_save
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-TEST_CASE_0 = [  # single channel 2D, batch 16, no residual
+TEST_CASE_0 = [
     {
         "spatial_dims": 2,
         "in_channels": 1,
@@ -36,7 +35,7 @@
     (16, 3, 32, 32),
 ]
 
-TEST_CASE_1 = [  # single channel 2D, batch 16
+TEST_CASE_1 = [
     {
         "spatial_dims": 2,
         "in_channels": 1,
@@ -49,7 +48,7 @@
     (16, 3, 32, 32),
 ]
 
-TEST_CASE_2 = [  # single channel 3D, batch 16
+TEST_CASE_2 = [
     {
         "spatial_dims": 3,
         "in_channels": 1,
@@ -62,7 +61,7 @@
     (16, 3, 32, 24, 48),
 ]
 
-TEST_CASE_3 = [  # 4-channel 3D, batch 16
+TEST_CASE_3 = [
     {
         "spatial_dims": 3,
         "in_channels": 4,
@@ -75,93 +74,7 @@
     (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]
-
-ILL_CASES = [
-    [
-        {  # len(channels) < 2
-            "spatial_dims": 2,
-            "in_channels": 1,
-            "out_channels": 3,
-            "channels": (16,),
-            "strides": (2, 2),
-            "num_res_units": 0,
-        }
-    ],
-    [
-        {  # len(strides) < len(channels) - 1
-            "spatial_dims": 2,
-            "in_channels": 1,
-            "out_channels": 3,
-            "channels": (8, 8, 8),
-            "strides": (2,),
-            "num_res_units": 0,
-        }
-    ],
-    [
-        {  # len(kernel_size) = 3, spatial_dims = 2
-            "spatial_dims": 2,
-            "in_channels": 1,
-            "out_channels": 3,
-            "channels": (8, 8, 8),
-            "strides": (2, 2),
-            "kernel_size": (3, 3, 3),
-        }
-    ],
-    [
-        {  # len(up_kernel_size) = 2, spatial_dims = 3
-            "spatial_dims": 3,
-            "in_channels": 1,
-            "out_channels": 3,
-            "channels": (8, 8, 8),
-            "strides": (2, 2),
-            "up_kernel_size": (3, 3),
-        }
-    ],
-]
+CASES = [TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]
 
 
 class TestCheckpointUNet(unittest.TestCase):
@@ -179,29 +92,51 @@ def test_script(self):
         test_data = torch.randn(16, 1, 32, 32)
         test_script_save(net, test_data)
 
-    def test_script_without_running_stats(self):
-        net = CheckpointUNet(
-            spatial_dims=2,
-            in_channels=1,
-            out_channels=3,
-            channels=(16, 32, 64),
-            strides=(2, 2),
-            num_res_units=0,
-            norm=("batch", {"track_running_stats": False}),
-        )
-        test_data = torch.randn(16, 1, 16, 4)
-        test_script_save(net, test_data)
-
     def test_ill_input_shape(self):
         net = CheckpointUNet(spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2))
         with eval_mode(net):
             with self.assertRaisesRegex(RuntimeError, "Sizes of tensors must match"):
                 net.forward(torch.randn(2, 1, 16, 5))
 
-    @parameterized.expand(ILL_CASES)
-    def test_ill_input_hyper_params(self, input_param):
-        with self.assertRaises(ValueError):
-            _ = CheckpointUNet(**input_param)
+    def test_checkpointing_equivalence_eval(self):
+        """Ensure that CheckpointUNet matches standard UNet in eval mode (checkpointing 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)
+
+        net_ckpt = CheckpointUNet(**params).to(device)
+        net_plain = UNet(**params).to(device)
+
+        with eval_mode(net_ckpt), eval_mode(net_plain):
+            y_ckpt = net_ckpt(x)
+            y_plain = net_plain(x)
+
+        # checkpointing should not change outputs in eval mode
+        self.assertTrue(torch.allclose(y_ckpt, y_plain, atol=1e-6, rtol=1e-5))
+
+    def test_checkpointing_activates_training(self):
+        """Ensure checkpointing triggers recomputation under training and gradients propagate."""
+        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)
+
+        # checkpointing should reduce activation memory use; we can't directly assert memory savings
+        # but we can confirm no runtime errors and gradients propagate correctly
+        self.assertIsNotNone(grad_norm)
 
 
 if __name__ == "__main__":
