Skip to content

feat: add activation checkpointing to unet #8554

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 13 commits into
base: dev
Choose a base branch
from
+58 −1
Open

feat: add activation checkpointing to unet #8554

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
de2b6bd
feat: add optional gradient checkpointing to unet
Sep 3, 2025
66edcb5
fix: small ruff issue
Sep 3, 2025
e66e357
Update monai/networks/nets/unet.py
ferreirafabio80 Sep 4, 2025
feefcaa
docs: update docstrings
Sep 4, 2025
e112457
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 4, 2025
f673ca1
fix: avoid BatchNorm subblocks
Sep 4, 2025
69540ff
fix: revert batch norm changes
Sep 4, 2025
42ec757
refactor: creates a subclass of UNet and overrides the get connection…
Oct 1, 2025
a2e8474
chore: remove use checkpointing from doc string
Oct 1, 2025
4c4782e
fix: linting issues
Oct 2, 2025
515c659
feat: add activation checkpointing to down and up paths to be more ef…
Oct 8, 2025
da5a3a4
refactor: move activation checkpointing wrapper to blocks
Nov 4, 2025
43dec88
chore: add docstrings to checkpointed unet
Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions monai/networks/blocks/activation_checkpointing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))
Comment on lines +32 to +41
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Gate checkpointing to active training passes.
The docstring promises training-only checkpointing, but forward always recomputes, so eval/no-grad still pays the checkpoint dispatch. Wrap the call with self.training, torch.is_grad_enabled(), and an x.requires_grad check, falling back to the plain module call otherwise, to avoid needless recompute overhead while preserving the memory trade-off during training.(docs.pytorch.org)

     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))
+        if self.training and torch.is_grad_enabled() and x.requires_grad:
+            return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+        return cast(torch.Tensor, self.module(x))
🤖 Prompt for AI Agents 
In monai/networks/blocks/activation_checkpointing.py around lines 32 to 41,
forward always calls checkpoint(self.module, x, use_reentrant=False) even during
eval/no-grad, causing unnecessary recompute; change it to only use
torch.utils.checkpoint when running training and gradients are enabled: check
self.training and torch.is_grad_enabled() and that input tensor x.requires_grad
before calling checkpoint(..., use_reentrant=False); otherwise call and return
self.module(x) directly to avoid unnecessary checkpoint overhead while
preserving training memory savings.

18 changes: 17 additions & 1 deletion monai/networks/nets/unet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -298,4 +299,19 @@ 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.
"""
Comment on lines +302 to +308
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Replace EN DASH with regular hyphen.

Line 306 contains – (EN DASH) instead of - (HYPHEN-MINUS).

-    intermediate activations inside encoder–decoder connections are recomputed in
+    intermediate activations inside encoder-decoder connections are recomputed in
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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.
"""
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.
"""
🧰 Tools
🪛 Ruff (0.14.3)

306-306: Docstring contains ambiguous – (EN DASH). Did you mean - (HYPHEN-MINUS)?

(RUF002)

🤖 Prompt for AI Agents 
In monai/networks/nets/unet.py around lines 302 to 308, the class docstring
contains an EN DASH (`–`) on line 306 used in "encoder–decoder"; replace that
character with a regular hyphen-minus (`-`) so it reads "encoder-decoder".
Update the docstring text only, keeping punctuation and spacing consistent.


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)
return super()._get_connection_block(down_path, up_path, subblock)
Comment on lines +302 to +314
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

Verify unit test coverage for CheckpointUNet.

Past review comments requested tests that confirm: (a) forward parity with UNet in eval/no-grad mode, (b) checkpointing engages only during training.

🏁 Script executed:

#!/bin/bash
# Search for test files covering CheckpointUNet
rg -n --type=py -C3 'CheckpointUNet|CheckpointUnet|checkpointunet' -g '*test*.py'

Length of output: 0


---

🏁 Script executed:

```shell
#!/bin/bash
# Broader search for test files in common test directories
find . -type f -name "*test*.py" -path "*/test*" | head -20

Length of output: 1066

🏁 Script executed:

#!/bin/bash
# Search for CheckpointUNet anywhere in the codebase (not just test files)
rg -n 'CheckpointUNet' --type=py

Length of output: 135

🏁 Script executed:

#!/bin/bash
# Check for UNet tests
rg -n 'class.*UNet.*Test|def test.*unet' --type=py -i

Length of output: 3075


---

🏁 Script executed:

```shell
#!/bin/bash
# Inspect test_unet.py for any CheckpointUNet references
rg -n 'CheckpointUNet|checkpoint' tests/networks/nets/test_unet.py -i

Length of output: 0

Add unit tests for CheckpointUNet to validate wrapper behavior.

CheckpointUNet has zero test coverage. Tests must confirm: (a) forward pass produces identical results to parent UNet in eval mode or with torch.no_grad(), and (b) activation checkpointing engages only during training with gradients enabled.

🧰 Tools
🪛 Ruff (0.14.3)

306-306: Docstring contains ambiguous – (EN DASH). Did you mean - (HYPHEN-MINUS)?

(RUF002)

🤖 Prompt for AI Agents 
In monai/networks/nets/unet.py around lines 302 to 314, add unit tests (e.g.,
monai/networks/nets/test_unet_checkpoint.py) that create a UNet and a
CheckpointUNet with identical seeded initialization and the same input tensor;
assert their forward outputs are equal in eval mode and when wrapped with
torch.no_grad(); then verify activation checkpointing is active only during
training with gradients by monkeypatching or wrapping
ActivationCheckpointWrapper.forward to count invocations: run a training
forward+backward (output.sum().backward()) with requires_grad enabled and assert
the wrapper.forward is invoked more than once (indicating recomputation), and
run the same in eval or torch.no_grad() and assert it is invoked exactly once.
Ensure deterministic seeding, zero gradients between runs, and use
torch.allclose with a tight tolerance for output comparisons.

Comment on lines +310 to +314
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

Add Google-style docstring to overridden method.

Per coding guidelines, all definitions require docstrings with Args/Returns sections.

As per coding guidelines.

     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)
🤖 Prompt for AI Agents 
In monai/networks/nets/unet.py around lines 310 to 314, the overridden
_get_connection_block method is missing a Google-style docstring; add a
docstring immediately above the def that follows Google style with a short
summary line, an Args section documenting down_path (nn.Module): the down path
module, up_path (nn.Module): the up path module, and subblock (nn.Module): the
connecting subblock, and a Returns section documenting nn.Module: the connection
block returned (note the method wraps the three inputs with
ActivationCheckpointWrapper and delegates to super()._get_connection_block);
keep wording concise and include types for each parameter and the return.



Unet = UNet
Loading