Add ConvNeXt for 1D, 2D and 3D inputs

[VenkateswarluNagineni]

Fixes #4798.

Description

Adds ConvNeXt (A ConvNet for the 2020s) as a classification backbone. MONAI has no ConvNeXt today, even though MedNeXt derives from it — MedNeXt is a segmentation architecture, so the modern convolutional backbone is still missing alongside resnet/densenet/senet/efficientnet.

This follows the direction @wyli gave on the issue:

I think from monai's perspective we can add value to it by supporting both 2d and 3d, and making use of existing monai network layers/blocks whenever possible.

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 Linear layers 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 / Pool factories for all dimension-agnostic ops
• DropPath (monai.networks.layers) for stochastic depth — it is already shape-agnostic, and the DropPath(p) if p > 0 else nn.Identity() guard and torch.linspace schedule follow swin_unetr.py
• get_act_layer for the act argument
• trunc_normal_ for the reference's std=0.02 initialisation

The structure follows densenet.py (the features / class_layers split, spatial_dims, in_channels, out_channels ordering, OrderedDict module naming, real subclasses for variants, alias block).

LayerNormNd

ConvNeXt normalises over the channel dimension of a channels-first feature map. torch.nn.LayerNorm normalises over the trailing dimensions, so it expects channels-last and cannot be used directly, and MONAI has no channels-first equivalent (Norm.LAYER is a plain nn.LayerNorm alias). So this PR adds LayerNormNd, which normalises dimension 1 of a (batch, channel, *spatial) tensor for any number of spatial dimensions.

It computes the statistics explicitly over dim=1 instead of permuting to channels-last, which keeps it uniform across ranks and TorchScript-friendly. A test asserts it equals nn.LayerNorm applied to the equivalent permuted tensor.

I've kept it in convnext.py to keep this PR's surface small — happy to promote it to monai/networks/layers/ as a public layer if you'd prefer, since nothing else in MONAI currently offers this.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New feature (non-breaking change which adds functionality).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

Docs are updated (docs/source/networks.rst, following the DenseNet entry pattern — base class with :members:, variants without), but I haven't run make html locally, and I'm on Windows so I ran the test files directly with unittest and the linters individually rather than through runtests.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):

• Faithfulness to the reference. The default variants match the parameter counts published by facebookresearch/ConvNeXt exactly for ImageNet-1k (Tiny 28,589,128 / Small 50,223,688 / Base 88,591,464; Large 197,767,336 and XLarge 350,196,968 also verified locally). This checks the architecture is correct rather than merely shape-correct.
• Shapes for 1D, 2D (non-cubic input) and 3D, over all five variants, plus a case with stochastic depth and layer scale disabled.
• LayerNormNd equivalence to nn.LayerNorm on the permuted tensor, for 2D and 3D (max difference ~3e-07).
• TorchScript via test_script_save for every case, including both the layer-scale and no-layer-scale paths.
• Drop path schedule ascends linearly from 0 to drop_path_rate across the whole network.
• Ill-args raise ValueError.

Also ran tests/networks/nets/test_densenet.py and test_mednext.py (56 tests) to confirm the nets/__init__.py addition doesn't regress anything.

Checks run on Windows with torch 2.13.0+cpu: ruff check, black --check, isort --check-only all clean on the changed files; mypy monai/networks/nets/convnext.py reports no errors for this file. (I ran the test files directly with unittest rather than the full ./runtests.sh, which expects a POSIX shell.)

Notes

• No pretrained support in this PR. The reference 2D weights could be mapped onto this implementation in a follow-up (the 1x1-conv vs Linear pointwise weights differ only by a reshape), but that felt like a separate change; MedNeXt likewise ships without pretrained weights.
• Every spatial dimension of the input should be divisible by 32 (patchify stem of 4 plus three downsamplings), which is documented in the class docstring.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 04cb79cb-2c20-4d85-ae11-a6f932691363

📥 Commits

Reviewing files that changed from the base of the PR and between c4d4ecf and a66704b.

📒 Files selected for processing (2)

• monai/networks/nets/convnext.py
• tests/networks/nets/test_convnext.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tests/networks/nets/test_convnext.py
• monai/networks/nets/convnext.py

---

📝 Walkthrough

Walkthrough

Adds 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)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: adding ConvNeXt support for 1D, 2D, and 3D inputs.
Description check	✅ Passed	The description follows the template and covers the feature, tests, and doc updates, though some checklist items remain unchecked.
Linked Issues check	✅ Passed	The PR implements the requested ConvNeXt architecture for MONAI and aligns with issue #4798.
Out of Scope Changes check	✅ Passed	The changes stay focused on ConvNeXt implementation, exports, docs, and tests, with no clear unrelated additions.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

monai/networks/nets/convnext.py (1)

70-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Complete 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 win

Exercise the aliases claimed by this test matrix.

TEST_VARIANT_CASES contains only canonical ConvNeXt* 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

📥 Commits

Reviewing files that changed from the base of the PR and between 3a458fe and c4d4ecf.

📒 Files selected for processing (4)

• docs/source/networks.rst
• monai/networks/nets/__init__.py
• monai/networks/nets/convnext.py
• tests/networks/nets/test_convnext.py

[VenkateswarluNagineni]

Thanks for the review. Addressed in the latest commit:

• Input validation — added. I confirmed the kernel_size case is a real trap: an even kernel doesn't preserve the spatial size at stride 1, so the residual add fails with an opaque RuntimeError: The size of tensor a (8) must match tensor b (9) deep in forward. It's now rejected in __init__ with a clear message, along with empty/mismatched/non-positive depths/features and an out-of-range drop_path_rate. All are covered by test_ill_arg and documented in the class Raises.
• Alias coverage — fixed. TEST_VARIANT_CASES now mixes the exported lowercase/mixed-case aliases (convnext_tiny, Convnext_base, …) with the canonical classes, so a broken package export is caught too.

I've left the per-forward docstrings out on purpose: the peer classification nets in this package (densenet.py, senet.py) document behavior at the class level and don't docstring forward/private helpers, and the ConvNeXt/_ConvNeXtBlock/LayerNormNd class docstrings already cover the tensor in/out contract. Following the existing house style there rather than adding docstrings the neighbouring modules don't have.