Add ConvNeXt for 1D, 2D and 3D inputs#8995
Open
VenkateswarluNagineni wants to merge 1 commit into
Open
Conversation
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>
VenkateswarluNagineni
requested review from
KumoLiu,
Nic-Ma and
ericspod
as code owners
Contributor
📝 WalkthroughWalkthroughAdds a full ConvNeXt implementation with configurable spatial dimensions, staged blocks, stochastic depth, layer scaling, initialization, and preset model variants. Exports the models and aliases through the networks package, documents them in Sphinx, and adds tests covering shapes, scripting, reference parameter counts, normalization, scheduling, scaling, and invalid arguments. 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}
Contributor
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
Comment on lines
+174
to
+188
| depths: Sequence[int] = (3, 3, 9, 3), | ||
| features: Sequence[int] = (96, 192, 384, 768), | ||
| drop_path_rate: float = 0.0, | ||
| layer_scale_init_value: float = 1e-6, | ||
| kernel_size: int = 7, | ||
| act: str | tuple = "gelu", | ||
| ) -> None: | ||
| super().__init__() | ||
|
|
||
| if spatial_dims not in (1, 2, 3): | ||
| raise ValueError(f"`spatial_dims` should be 1, 2 or 3, got {spatial_dims}.") | ||
| if len(depths) != len(features): | ||
| raise ValueError( | ||
| f"`depths` and `features` should have the same length, got {len(depths)} and {len(features)}." | ||
| ) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Validate architecture parameters before constructing the network.
Equal-length empty sequences pass validation and then crash at features[0]. Even kernel_size values enlarge the depthwise-convolution output, breaking the residual addition. Also reject nonpositive depths/features and drop_path_rate outside [0, 1].
Proposed validation
- if len(depths) != len(features):
+ if len(depths) != 4 or len(features) != 4:
raise ValueError(
- f"`depths` and `features` should have the same length, got {len(depths)} and {len(features)}."
+ "`depths` and `features` must each describe four stages."
)
+ if any(depth <= 0 for depth in depths):
+ raise ValueError("Every value in `depths` must be positive.")
+ if any(feature <= 0 for feature in features):
+ raise ValueError("Every value in `features` must be positive.")
+ if not 0.0 <= drop_path_rate <= 1.0:
+ raise ValueError("`drop_path_rate` must be between 0 and 1.")
+ if kernel_size <= 0 or kernel_size % 2 == 0:
+ raise ValueError("`kernel_size` must be a positive odd integer.")Add corresponding invalid-argument tests.
As per path instructions, examine logical 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.
Suggested change
| depths: Sequence[int] = (3, 3, 9, 3), | |
| features: Sequence[int] = (96, 192, 384, 768), | |
| drop_path_rate: float = 0.0, | |
| layer_scale_init_value: float = 1e-6, | |
| kernel_size: int = 7, | |
| act: str | tuple = "gelu", | |
| ) -> None: | |
| super().__init__() | |
| if spatial_dims not in (1, 2, 3): | |
| raise ValueError(f"`spatial_dims` should be 1, 2 or 3, got {spatial_dims}.") | |
| if len(depths) != len(features): | |
| raise ValueError( | |
| f"`depths` and `features` should have the same length, got {len(depths)} and {len(features)}." | |
| ) | |
| depths: Sequence[int] = (3, 3, 9, 3), | |
| features: Sequence[int] = (96, 192, 384, 768), | |
| drop_path_rate: float = 0.0, | |
| layer_scale_init_value: float = 1e-6, | |
| kernel_size: int = 7, | |
| act: str | tuple = "gelu", | |
| ) -> None: | |
| super().__init__() | |
| if spatial_dims not in (1, 2, 3): | |
| raise ValueError(f"`spatial_dims` should be 1, 2 or 3, got {spatial_dims}.") | |
| if len(depths) != 4 or len(features) != 4: | |
| raise ValueError( | |
| "`depths` and `features` must each describe four stages." | |
| ) | |
| if any(depth <= 0 for depth in depths): | |
| raise ValueError("Every value in `depths` must be positive.") | |
| if any(feature <= 0 for feature in features): | |
| raise ValueError("Every value in `features` must be positive.") | |
| if not 0.0 <= drop_path_rate <= 1.0: | |
| raise ValueError("`drop_path_rate` must be between 0 and 1.") | |
| if kernel_size <= 0 or kernel_size % 2 == 0: | |
| raise ValueError("`kernel_size` must be a positive odd integer.") |
🤖 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 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.
Source: Path instructions
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.