refactor: move activation checkpointing wrapper to blocks

Author: Fabio Ferreira

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

@@ -13,30 +13,18 @@
 
 import warnings
 from collections.abc import Sequence
-from typing import cast
 
 import torch
 import torch.nn as nn
-from torch.utils.checkpoint import checkpoint
 
+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"]
 
 
-class _ActivationCheckpointWrapper(nn.Module):
-    """Apply activation checkpointing to the wrapped module during training."""
-
-    def __init__(self, module: nn.Module) -> None:
-        super().__init__()
-        self.module = module
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
-
-
 class UNet(nn.Module):
     """
     Enhanced version of UNet which has residual units implemented with the ResidualUnit class.
@@ -313,9 +301,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class CheckpointUNet(UNet):
     def _get_connection_block(self, down_path: nn.Module, up_path: nn.Module, subblock: nn.Module) -> nn.Module:
-        subblock = _ActivationCheckpointWrapper(subblock)
-        down_path = _ActivationCheckpointWrapper(down_path)
-        up_path = _ActivationCheckpointWrapper(up_path)
+        subblock = ActivationCheckpointWrapper(subblock)
+        down_path = ActivationCheckpointWrapper(down_path)
+        up_path = ActivationCheckpointWrapper(up_path)
         return super()._get_connection_block(down_path, up_path, subblock)