Add ConvNeXt for 1D, 2D and 3D inputs#8995
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
monai/networks/nets/convnext.py (1)
70-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete Google-style documentation for every new definition.
monai/networks/nets/convnext.py#L70-L74: document the input and returned tensor.monai/networks/nets/convnext.py#L120-L131: document the block input and residual output.monai/networks/nets/convnext.py#L250-L253: document classifier input and output.monai/networks/nets/convnext.py#L256-L363: document variant constructor parameters and exceptions.tests/networks/nets/test_convnext.py#L82-L92: document parameterized test arguments.tests/networks/nets/test_convnext.py#L147-L151: add a concise test docstring.As per path instructions, docstrings must be present for all definitions using appropriate Google-style sections.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/networks/nets/convnext.py` around lines 70 - 74, Complete the Google-style docstrings for all listed definitions: in monai/networks/nets/convnext.py ranges 70-74, 120-131, and 250-253, document tensor inputs and outputs; in the variant constructors at 256-363, document parameters and raised exceptions; in tests/networks/nets/test_convnext.py ranges 82-92 and 147-151, document parameterized arguments and add a concise test description. Ensure every Args, Returns, and Raises section accurately reflects the existing signatures and behavior.Source: Path instructions
tests/networks/nets/test_convnext.py (1)
66-70: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise the aliases claimed by this test matrix.
TEST_VARIANT_CASEScontains only canonicalConvNeXt*classes, despite the comment claiming alias coverage. Import and include the exported aliases so broken package exports are detected.As per path instructions, ensure new or modified definitions are covered by unit tests.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/networks/nets/test_convnext.py` around lines 66 - 70, Update TEST_VARIANT_CASES to include the exported ConvNeXt aliases alongside the canonical classes, preserving coverage across both TEST_CASE_1 and TEST_CASE_2. Import the aliases from the same package source and ensure each alias is exercised by the existing matrix so package export regressions are tested.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@monai/networks/nets/convnext.py`:
- Around line 174-188: Update ConvNeXt.__init__ validation to reject empty
depths/features, any nonpositive depth or feature value, nonpositive or even
kernel_size values that would break residual shape matching, and drop_path_rate
outside the inclusive [0, 1] range. Preserve the existing spatial_dims and
equal-length checks, and add unit tests covering each invalid-argument case.
---
Nitpick comments:
In `@monai/networks/nets/convnext.py`:
- Around line 70-74: Complete the Google-style docstrings for all listed
definitions: in monai/networks/nets/convnext.py ranges 70-74, 120-131, and
250-253, document tensor inputs and outputs; in the variant constructors at
256-363, document parameters and raised exceptions; in
tests/networks/nets/test_convnext.py ranges 82-92 and 147-151, document
parameterized arguments and add a concise test description. Ensure every Args,
Returns, and Raises section accurately reflects the existing signatures and
behavior.
In `@tests/networks/nets/test_convnext.py`:
- Around line 66-70: Update TEST_VARIANT_CASES to include the exported ConvNeXt
aliases alongside the canonical classes, preserving coverage across both
TEST_CASE_1 and TEST_CASE_2. Import the aliases from the same package source and
ensure each alias is exercised by the existing matrix so package export
regressions are tested.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: dfedbdb4-57cf-4da5-8fb8-d068ea0cbb0d
📒 Files selected for processing (4)
docs/source/networks.rstmonai/networks/nets/__init__.pymonai/networks/nets/convnext.pytests/networks/nets/test_convnext.py
|
Thanks for the review. Addressed in the latest commit:
I've left the per- |
Would definitely do this |
|
Out of curiosity why not Port the V2 version instead? (Cool contribution!) |
|
Thanks @aymuos15! Public ConvNeXt V2 — good call, and I'd like to add it, but as a focused follow-up rather than swapping it in here:
Happy to fold V1 + V2 into this single PR instead if you'd prefer — just say the word. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
monai/networks/layers/layer_norm_nd.py (1)
42-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument
forward’s tensor contract.Add a Google-style docstring for
xand the normalized return tensor.As per path instructions, “Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/networks/layers/layer_norm_nd.py` around lines 42 - 46, Add a Google-style docstring to LayerNormND.forward documenting the x input tensor and the normalized tensor returned, including their expected shape or dimensionality where established by the class contract. Keep the normalization implementation unchanged and document any raised exceptions only if this method explicitly raises them.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@monai/networks/layers/layer_norm_nd.py`:
- Around line 33-40: Validate num_channels and spatial_dims at the start of
LayerNormND.__init__, rejecting zero or negative channel counts and negative
spatial dimensions before creating parameters or param_shape; preserve
valid-dimension behavior and add unit tests covering each invalid argument.
---
Nitpick comments:
In `@monai/networks/layers/layer_norm_nd.py`:
- Around line 42-46: Add a Google-style docstring to LayerNormND.forward
documenting the x input tensor and the normalized tensor returned, including
their expected shape or dimensionality where established by the class contract.
Keep the normalization implementation unchanged and document any raised
exceptions only if this method explicitly raises them.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 47c89d08-0c0f-492d-9b73-fbbca5094a30
📒 Files selected for processing (6)
docs/source/networks.rstmonai/networks/layers/__init__.pymonai/networks/layers/layer_norm_nd.pymonai/networks/nets/convnext.pytests/networks/layers/test_layer_norm_nd.pytests/networks/nets/test_convnext.py
💤 Files with no reviewable changes (1)
- tests/networks/nets/test_convnext.py
| def __init__(self, num_channels: int, spatial_dims: int, eps: float = 1e-6) -> None: | ||
| super().__init__() | ||
| self.eps = eps | ||
| self.weight = nn.Parameter(torch.ones(num_channels)) | ||
| self.bias = nn.Parameter(torch.zeros(num_channels)) | ||
| # broadcast the affine parameters against (batch, channel, *spatial); precomputed so that the | ||
| # module is scriptable without inspecting the rank of the input at runtime. | ||
| self.param_shape = [1, num_channels] + [1] * spatial_dims |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject invalid layer dimensions.
Negative spatial_dims produces an affine shape that can broadcast over the wrong axis; num_channels=0 also creates an unusable layer. Validate both and add invalid-argument tests.
Proposed fix
def __init__(self, num_channels: int, spatial_dims: int, eps: float = 1e-6) -> None:
super().__init__()
+ if num_channels <= 0:
+ raise ValueError("num_channels must be positive.")
+ if spatial_dims < 0:
+ raise ValueError("spatial_dims must be non-negative.")
self.eps = epsAs per path instructions, “Examine code for logical error or inconsistencies” and ensure modified definitions have unit tests.
📝 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.
| def __init__(self, num_channels: int, spatial_dims: int, eps: float = 1e-6) -> None: | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(num_channels)) | |
| self.bias = nn.Parameter(torch.zeros(num_channels)) | |
| # broadcast the affine parameters against (batch, channel, *spatial); precomputed so that the | |
| # module is scriptable without inspecting the rank of the input at runtime. | |
| self.param_shape = [1, num_channels] + [1] * spatial_dims | |
| def __init__(self, num_channels: int, spatial_dims: int, eps: float = 1e-6) -> None: | |
| super().__init__() | |
| if num_channels <= 0: | |
| raise ValueError("num_channels must be positive.") | |
| if spatial_dims < 0: | |
| raise ValueError("spatial_dims must be non-negative.") | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(num_channels)) | |
| self.bias = nn.Parameter(torch.zeros(num_channels)) | |
| # broadcast the affine parameters against (batch, channel, *spatial); precomputed so that the | |
| # module is scriptable without inspecting the rank of the input at runtime. | |
| self.param_shape = [1, num_channels] + [1] * spatial_dims |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@monai/networks/layers/layer_norm_nd.py` around lines 33 - 40, Validate
num_channels and spatial_dims at the start of LayerNormND.__init__, rejecting
zero or negative channel counts and negative spatial dimensions before creating
parameters or param_shape; preserve valid-dimension behavior and add unit tests
covering each invalid argument.
Source: Path instructions
1c0b70c to
ca8fcb2
Compare
ConvNeXt is a modern convolutional classification backbone, and MONAI has no implementation of it although MedNeXt derives from it. Add one that supports volumetric inputs rather than only the 2D images of the reference implementation, so that it can be used as a backbone for medical images. The block is built from the existing dimension agnostic MONAI layers: the `Conv` and `Pool` factories, `DropPath` for stochastic depth, `get_act_layer` for the activation, and `trunc_normal_` for the reference initialisation. Channel normalisation needs a channels-first `LayerNorm`, which MONAI does not have: `torch.nn.LayerNorm` normalises over the trailing dimensions and so expects a channels-last layout. `LayerNormNd` normalises the channel dimension of a (batch, channel, *spatial) tensor for any number of spatial dimensions, and is verified against `torch.nn.LayerNorm` applied to the equivalent permuted tensor. The default variants match the parameter counts published for the reference 2D models, which is asserted by the tests. Signed-off-by: VenkateswarluNagineni <venkates2002@tamu.edu>
…ests Reject argument combinations that would otherwise fail deep in the forward pass with an opaque error: an even `kernel_size` cannot preserve the spatial size at stride 1 and so breaks the residual connection, and empty, mismatched or non-positive `depths`/`features` and an out-of-range `drop_path_rate` are now caught in `__init__` with a clear message. Document these in the class `Raises`. Also mix the exported lowercase aliases into the variant test matrix so that a broken package export is caught as well as a broken variant. Signed-off-by: VenkateswarluNagineni <venkates2002@tamu.edu>
for more information, see https://pre-commit.ci
Move the channels-first LayerNormNd out of the ConvNeXt net into monai/networks/layers so any channels-first backbone can reuse it, mirroring the DropPath layer. Export it from monai.networks.layers, document it in the Layers docs, and add a dedicated layer test that checks equivalence to nn.LayerNorm on the permuted channels-last tensor. Signed-off-by: VenkateswarluNagineni <venkates2002@tamu.edu>
Cast the bias to a tensor before `nn.init.constant_` (matching the densenet/resnet/hovernet weight-init idiom) so mypy no longer flags the `Tensor | Module` type, and document `LayerNormNd` under the `monai.networks.layers` current module so autodoc can import it. Signed-off-by: VenkateswarluNagineni <venkates2002@tamu.edu>
The channels-first and channels-last shapes in the class docstring were written as bare tuples, so docutils read the leading asterisk of `*spatial` as the start of inline emphasis and the docs build failed on the resulting warning. Mark them up as inline literals instead. Signed-off-by: VenkateswarluNagineni <venkates2002@tamu.edu>
ca8fcb2 to
998e5c1
Compare
Wasnt aware pytorch doesnt have the official weights for it. My bad. This LGTM |
2 participants
Fixes #4798.
Description
Adds ConvNeXt (
A ConvNet for the 2020s) as a classification backbone. MONAI has no ConvNeXt today, even thoughMedNeXtderives from it — MedNeXt is a segmentation architecture, so the modern convolutional backbone is still missing alongsideresnet/densenet/senet/efficientnet.This follows the direction @wyli gave on the issue:
Supports 1D, 2D and 3D. The reference implementation is 2D-only; every layer here is built through MONAI's dimension-agnostic factories, so volumetric inputs work with the same code path. The reference applies the pointwise convolutions as
Linearlayers on a permuted channels-last tensor; using 1x1 convolutions instead is mathematically equivalent and keeps the block agnostic to the number of spatial dimensions.Reuses existing MONAI layers rather than re-implementing them:
Conv/Poolfactories for all dimension-agnostic opsDropPath(monai.networks.layers) for stochastic depth — it is already shape-agnostic, and theDropPath(p) if p > 0 else nn.Identity()guard andtorch.linspaceschedule followswin_unetr.pyget_act_layerfor theactargumenttrunc_normal_for the reference'sstd=0.02initialisationThe structure follows
densenet.py(thefeatures/class_layerssplit,spatial_dims, in_channels, out_channelsordering,OrderedDictmodule naming, real subclasses for variants, alias block).LayerNormNdConvNeXt normalises over the channel dimension of a channels-first feature map.
torch.nn.LayerNormnormalises over the trailing dimensions, so it expects channels-last and cannot be used directly, and MONAI has no channels-first equivalent (Norm.LAYERis a plainnn.LayerNormalias). So this PR addsLayerNormNd, which normalises dimension 1 of a(batch, channel, *spatial)tensor for any number of spatial dimensions.It computes the statistics explicitly over
dim=1instead of permuting to channels-last, which keeps it uniform across ranks and TorchScript-friendly. A test asserts it equalsnn.LayerNormapplied to the equivalent permuted tensor.I've kept it in
convnext.pyto keep this PR's surface small — happy to promote it tomonai/networks/layers/as a public layer if you'd prefer, since nothing else in MONAI currently offers this.Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.Docs are updated (
docs/source/networks.rst, following theDenseNetentry pattern — base class with:members:, variants without), but I haven't runmake htmllocally, and I'm on Windows so I ran the test files directly withunittestand the linters individually rather than throughruntests.sh. Details below — happy to run anything else you'd like to see.Testing
New file
tests/networks/nets/test_convnext.py, 26 tests, all passing (python -m unittest tests.networks.nets.test_convnext):LayerNormNdequivalence tonn.LayerNormon the permuted tensor, for 2D and 3D (max difference ~3e-07).test_script_savefor every case, including both the layer-scale and no-layer-scale paths.drop_path_rateacross the whole network.ValueError.Also ran
tests/networks/nets/test_densenet.pyandtest_mednext.py(56 tests) to confirm thenets/__init__.pyaddition doesn't regress anything.Checks run on Windows with torch 2.13.0+cpu:
ruff check,black --check,isort --check-onlyall clean on the changed files;mypy monai/networks/nets/convnext.pyreports no errors for this file. (I ran the test files directly withunittestrather than the full./runtests.sh, which expects a POSIX shell.)Notes
pretrainedsupport in this PR. The reference 2D weights could be mapped onto this implementation in a follow-up (the 1x1-conv vsLinearpointwise weights differ only by a reshape), but that felt like a separate change;MedNeXtlikewise ships without pretrained weights.