From 461afd01469e1997580db6205a5020bdcfb70634 Mon Sep 17 00:00:00 2001 From: VenkateswarluNagineni Date: Wed, 15 Jul 2026 22:24:04 -0500 Subject: [PATCH 1/6] feat(networks): add ConvNeXt for 1D, 2D and 3D inputs 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 --- docs/source/networks.rst | 25 ++ monai/networks/nets/__init__.py | 19 ++ monai/networks/nets/convnext.py | 371 +++++++++++++++++++++++++++ tests/networks/nets/test_convnext.py | 155 +++++++++++ 4 files changed, 570 insertions(+) create mode 100644 monai/networks/nets/convnext.py create mode 100644 tests/networks/nets/test_convnext.py diff --git a/docs/source/networks.rst b/docs/source/networks.rst index e7709678b7..17f2ddc273 100644 --- a/docs/source/networks.rst +++ b/docs/source/networks.rst @@ -467,6 +467,31 @@ Nets .. autoclass:: AHNet :members: +`ConvNeXt` +~~~~~~~~~~ +.. autoclass:: ConvNeXt + :members: + +`ConvNeXtTiny` +~~~~~~~~~~~~~~ +.. autoclass:: ConvNeXtTiny + +`ConvNeXtSmall` +~~~~~~~~~~~~~~~ +.. autoclass:: ConvNeXtSmall + +`ConvNeXtBase` +~~~~~~~~~~~~~~ +.. autoclass:: ConvNeXtBase + +`ConvNeXtLarge` +~~~~~~~~~~~~~~~ +.. autoclass:: ConvNeXtLarge + +`ConvNeXtXLarge` +~~~~~~~~~~~~~~~~ +.. autoclass:: ConvNeXtXLarge + `DenseNet` ~~~~~~~~~~ .. autoclass:: DenseNet diff --git a/monai/networks/nets/__init__.py b/monai/networks/nets/__init__.py index fc0b33a0f0..8f8849eeaa 100644 --- a/monai/networks/nets/__init__.py +++ b/monai/networks/nets/__init__.py @@ -19,6 +19,25 @@ from .basic_unetplusplus import BasicUNetPlusPlus, BasicUnetPlusPlus, BasicunetPlusPlus, basicunetplusplus from .classifier import Classifier, Critic, Discriminator from .controlnet import ControlNet +from .convnext import ( + ConvNeXt, + Convnext, + Convnext_base, + Convnext_large, + Convnext_small, + Convnext_tiny, + Convnext_xlarge, + ConvNeXtBase, + ConvNeXtLarge, + ConvNeXtSmall, + ConvNeXtTiny, + ConvNeXtXLarge, + convnext_base, + convnext_large, + convnext_small, + convnext_tiny, + convnext_xlarge, +) from .daf3d import DAF3D from .densenet import ( DenseNet, diff --git a/monai/networks/nets/convnext.py b/monai/networks/nets/convnext.py new file mode 100644 index 0000000000..303b6e0539 --- /dev/null +++ b/monai/networks/nets/convnext.py @@ -0,0 +1,371 @@ +# 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 collections import OrderedDict +from collections.abc import Sequence + +import torch +import torch.nn as nn + +from monai.networks.layers.drop_path import DropPath +from monai.networks.layers.factories import Conv, Pool +from monai.networks.layers.utils import get_act_layer +from monai.networks.layers.weight_init import trunc_normal_ + +__all__ = [ + "ConvNeXt", + "Convnext", + "ConvNeXtTiny", + "Convnext_tiny", + "convnext_tiny", + "ConvNeXtSmall", + "Convnext_small", + "convnext_small", + "ConvNeXtBase", + "Convnext_base", + "convnext_base", + "ConvNeXtLarge", + "Convnext_large", + "convnext_large", + "ConvNeXtXLarge", + "Convnext_xlarge", + "convnext_xlarge", +] + + +class LayerNormNd(nn.Module): + """ + Layer normalization over the channel dimension of a channels-first tensor. + + `torch.nn.LayerNorm` normalizes over the trailing dimensions, so it expects a channels-last layout + such as (batch, *spatial, channel). Convolutional feature maps in MONAI are channels-first, + (batch, channel, *spatial), and this module normalizes those over the channel dimension only, + for any number of spatial dimensions. + + Args: + num_channels: number of channels of the input, i.e. the size of dimension 1. + spatial_dims: number of spatial dimensions of the input image. + eps: value added to the denominator for numerical stability. + """ + + 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 forward(self, x: torch.Tensor) -> torch.Tensor: + mean = x.mean(dim=1, keepdim=True) + var = (x - mean).pow(2).mean(dim=1, keepdim=True) + x = (x - mean) / torch.sqrt(var + self.eps) + return self.weight.view(self.param_shape) * x + self.bias.view(self.param_shape) + + +class _ConvNeXtBlock(nn.Module): + """ + A single ConvNeXt block: depthwise convolution, normalization, and an inverted bottleneck of two + pointwise convolutions, added back to the input through an optional layer scale and drop path. + + The reference implementation applies the pointwise convolutions as linear layers on a permuted, + channels-last tensor. Using 1x1 convolutions instead is equivalent, and keeps the block agnostic to + the number of spatial dimensions. + + Args: + spatial_dims: number of spatial dimensions of the input image. + dim: number of channels of the input and of the output. + kernel_size: size of the depthwise convolution kernel. + drop_path: stochastic depth rate. + layer_scale_init_value: initial value of the layer scale applied to the residual branch, + no layer scale is applied when this is not positive. + act: activation type and arguments. + """ + + def __init__( + self, + spatial_dims: int, + dim: int, + kernel_size: int = 7, + drop_path: float = 0.0, + layer_scale_init_value: float = 1e-6, + act: str | tuple = "gelu", + ) -> None: + super().__init__() + conv_type: type[nn.Conv1d | nn.Conv2d | nn.Conv3d] = Conv[Conv.CONV, spatial_dims] + + self.dwconv = conv_type(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim) + self.norm = LayerNormNd(dim, spatial_dims=spatial_dims) + self.pwconv1 = conv_type(dim, 4 * dim, kernel_size=1) + self.act = get_act_layer(name=act) + self.pwconv2 = conv_type(4 * dim, dim, kernel_size=1) + if layer_scale_init_value > 0: + self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim)) + else: + self.register_parameter("gamma", None) + self.param_shape = [1, dim] + [1] * spatial_dims + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + x = self.dwconv(x) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + gamma = self.gamma + if gamma is not None: + x = gamma.view(self.param_shape) * x + residual: torch.Tensor = self.drop_path(x) + return identity + residual + + +class ConvNeXt(nn.Module): + """ + ConvNeXt based on: `A ConvNet for the 2020s `_. + Adapted from the 2D reference implementation: https://github.com/facebookresearch/ConvNeXt. + + The network is a classification backbone built from four stages of :py:class:`_ConvNeXtBlock`, + separated by downsampling layers. Unlike the reference implementation it supports 1D, 2D and 3D + inputs, which makes it usable for volumetric medical images. + + Each downsampling layer halves the spatial size and the patchify stem reduces it by a factor of 4, + so every spatial dimension of the input should be divisible by 32. + + Args: + spatial_dims: number of spatial dimensions of the input image. + in_channels: number of the input channel. + out_channels: number of the output classes. + depths: number of blocks in each of the four stages. + features: number of channels in each of the four stages. + drop_path_rate: stochastic depth rate, increased linearly over the blocks. + layer_scale_init_value: initial value of the layer scale applied to each residual branch, + no layer scale is applied when this is not positive. + kernel_size: size of the depthwise convolution kernel of each block. + act: activation type and arguments. Defaults to gelu. + + Raises: + ValueError: when `spatial_dims` is not one of (1, 2, 3). + ValueError: when `depths` and `features` have different lengths. + + Example:: + + # 3D ConvNeXt-Tiny for binary classification of single channel volumes + net = ConvNeXtTiny(spatial_dims=3, in_channels=1, out_channels=2) + output = net(torch.randn(2, 1, 32, 32, 32)) # (2, 2) + """ + + def __init__( + self, + spatial_dims: int, + in_channels: int, + out_channels: int, + 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)}." + ) + + conv_type: type[nn.Conv1d | nn.Conv2d | nn.Conv3d] = Conv[Conv.CONV, spatial_dims] + avg_pool_type: type[nn.AdaptiveAvgPool1d | nn.AdaptiveAvgPool2d | nn.AdaptiveAvgPool3d] = Pool[ + Pool.ADAPTIVEAVG, spatial_dims + ] + + # stochastic depth increases linearly over the blocks of the whole network + drop_path_rates = [float(x) for x in torch.linspace(0, drop_path_rate, sum(depths))] + + self.features = nn.Sequential( + OrderedDict( + [ + ( + "stem", + nn.Sequential( + conv_type(in_channels, features[0], kernel_size=4, stride=4), + LayerNormNd(features[0], spatial_dims=spatial_dims), + ), + ) + ] + ) + ) + for i, depth in enumerate(depths): + if i > 0: + downsample = nn.Sequential( + LayerNormNd(features[i - 1], spatial_dims=spatial_dims), + conv_type(features[i - 1], features[i], kernel_size=2, stride=2), + ) + self.features.add_module(f"downsample{i}", downsample) + blocks = [ + _ConvNeXtBlock( + spatial_dims=spatial_dims, + dim=features[i], + kernel_size=kernel_size, + drop_path=drop_path_rates[sum(depths[:i]) + j], + layer_scale_init_value=layer_scale_init_value, + act=act, + ) + for j in range(depth) + ] + self.features.add_module(f"stage{i + 1}", nn.Sequential(*blocks)) + + # pooling and classification, the final normalization is over the pooled feature vector and so + # it is channels-last and uses `nn.LayerNorm` directly. + self.class_layers = nn.Sequential( + OrderedDict( + [ + ("pool", avg_pool_type(1)), + ("flatten", nn.Flatten(1)), + ("norm", nn.LayerNorm(features[-1], eps=1e-6)), + ("out", nn.Linear(features[-1], out_channels)), + ] + ) + ) + + for m in self.modules(): + if isinstance(m, (conv_type, nn.Linear)): + trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.features(x) + x = self.class_layers(x) + return x + + +class ConvNeXtTiny(ConvNeXt): + """ConvNeXt-T, the tiny variant of :py:class:`ConvNeXt`.""" + + def __init__( + self, + spatial_dims: int, + in_channels: int, + out_channels: int, + depths: Sequence[int] = (3, 3, 9, 3), + features: Sequence[int] = (96, 192, 384, 768), + **kwargs, + ) -> None: + super().__init__( + spatial_dims=spatial_dims, + in_channels=in_channels, + out_channels=out_channels, + depths=depths, + features=features, + **kwargs, + ) + + +class ConvNeXtSmall(ConvNeXt): + """ConvNeXt-S, the small variant of :py:class:`ConvNeXt`.""" + + def __init__( + self, + spatial_dims: int, + in_channels: int, + out_channels: int, + depths: Sequence[int] = (3, 3, 27, 3), + features: Sequence[int] = (96, 192, 384, 768), + **kwargs, + ) -> None: + super().__init__( + spatial_dims=spatial_dims, + in_channels=in_channels, + out_channels=out_channels, + depths=depths, + features=features, + **kwargs, + ) + + +class ConvNeXtBase(ConvNeXt): + """ConvNeXt-B, the base variant of :py:class:`ConvNeXt`.""" + + def __init__( + self, + spatial_dims: int, + in_channels: int, + out_channels: int, + depths: Sequence[int] = (3, 3, 27, 3), + features: Sequence[int] = (128, 256, 512, 1024), + **kwargs, + ) -> None: + super().__init__( + spatial_dims=spatial_dims, + in_channels=in_channels, + out_channels=out_channels, + depths=depths, + features=features, + **kwargs, + ) + + +class ConvNeXtLarge(ConvNeXt): + """ConvNeXt-L, the large variant of :py:class:`ConvNeXt`.""" + + def __init__( + self, + spatial_dims: int, + in_channels: int, + out_channels: int, + depths: Sequence[int] = (3, 3, 27, 3), + features: Sequence[int] = (192, 384, 768, 1536), + **kwargs, + ) -> None: + super().__init__( + spatial_dims=spatial_dims, + in_channels=in_channels, + out_channels=out_channels, + depths=depths, + features=features, + **kwargs, + ) + + +class ConvNeXtXLarge(ConvNeXt): + """ConvNeXt-XL, the extra large variant of :py:class:`ConvNeXt`.""" + + def __init__( + self, + spatial_dims: int, + in_channels: int, + out_channels: int, + depths: Sequence[int] = (3, 3, 27, 3), + features: Sequence[int] = (256, 512, 1024, 2048), + **kwargs, + ) -> None: + super().__init__( + spatial_dims=spatial_dims, + in_channels=in_channels, + out_channels=out_channels, + depths=depths, + features=features, + **kwargs, + ) + + +Convnext = ConvNeXt +Convnext_tiny = convnext_tiny = ConvNeXtTiny +Convnext_small = convnext_small = ConvNeXtSmall +Convnext_base = convnext_base = ConvNeXtBase +Convnext_large = convnext_large = ConvNeXtLarge +Convnext_xlarge = convnext_xlarge = ConvNeXtXLarge diff --git a/tests/networks/nets/test_convnext.py b/tests/networks/nets/test_convnext.py new file mode 100644 index 0000000000..eada91b356 --- /dev/null +++ b/tests/networks/nets/test_convnext.py @@ -0,0 +1,155 @@ +# 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 + +import unittest + +import torch +import torch.nn as nn +from parameterized import parameterized + +from monai.networks import eval_mode +from monai.networks.nets import ConvNeXt, ConvNeXtBase, ConvNeXtLarge, ConvNeXtSmall, ConvNeXtTiny, ConvNeXtXLarge +from monai.networks.nets.convnext import LayerNormNd +from tests.test_utils import skip_if_quick, test_script_save + +device = "cuda" if torch.cuda.is_available() else "cpu" + +# the variants only differ in `depths` and `features`, `features` is overridden throughout to keep the +# tests small while still exercising each variant's real depth configuration. +SMALL_FEATURES = (4, 8, 16, 32) + +TEST_CASE_1 = [ # 2-channel 3D, batch 2 + {"spatial_dims": 3, "in_channels": 2, "out_channels": 3, "features": SMALL_FEATURES}, + (2, 2, 32, 32, 32), + (2, 3), +] +TEST_CASE_2 = [ # 2-channel 2D, batch 2, non-cubic input + {"spatial_dims": 2, "in_channels": 2, "out_channels": 3, "features": SMALL_FEATURES, "act": "relu"}, + (2, 2, 32, 64), + (2, 3), +] +TEST_CASE_3 = [ # 1-channel 1D, batch 1 + {"spatial_dims": 1, "in_channels": 1, "out_channels": 3, "features": SMALL_FEATURES}, + (1, 1, 64), + (1, 3), +] +TEST_CASE_4 = [ # stochastic depth and no layer scale + { + "spatial_dims": 3, + "in_channels": 1, + "out_channels": 2, + "features": SMALL_FEATURES, + "drop_path_rate": 0.2, + "layer_scale_init_value": 0.0, + "kernel_size": 3, + }, + (2, 1, 32, 32, 32), + (2, 2), +] + +CASES = [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4] + +TEST_CASES = [] +for case in CASES: + TEST_CASES.append([ConvNeXt, *case]) + +# every variant, over the 2D and 3D cases, exercising the aliases alongside the canonical names +TEST_VARIANT_CASES = [] +for model in [ConvNeXtTiny, ConvNeXtSmall, ConvNeXtBase, ConvNeXtLarge, ConvNeXtXLarge]: + for case in [TEST_CASE_1, TEST_CASE_2]: + TEST_VARIANT_CASES.append([model, *case]) + +TEST_SCRIPT_CASES = [[ConvNeXt, *case] for case in CASES] + +# published ImageNet-1k parameter counts of the official 2D implementation, +# https://github.com/facebookresearch/ConvNeXt +TEST_REFERENCE_PARAMS = [[ConvNeXtTiny, 28_589_128], [ConvNeXtSmall, 50_223_688], [ConvNeXtBase, 88_591_464]] + + +class TestConvNeXt(unittest.TestCase): + + @parameterized.expand(TEST_CASES + TEST_VARIANT_CASES) + def test_convnext_shape(self, model, input_param, input_shape, expected_shape): + net = model(**input_param).to(device) + with eval_mode(net): + result = net.forward(torch.randn(input_shape).to(device)) + self.assertEqual(result.shape, expected_shape) + + @parameterized.expand(TEST_SCRIPT_CASES) + def test_script(self, model, input_param, input_shape, expected_shape): + net = model(**input_param) + test_data = torch.randn(input_shape) + test_script_save(net, test_data) + + @parameterized.expand(TEST_REFERENCE_PARAMS) + @skip_if_quick + def test_reference_parameter_count(self, model, expected_params): + """The default variants should match the published parameter counts of the 2D reference models.""" + net = model(spatial_dims=2, in_channels=3, out_channels=1000) + self.assertEqual(sum(p.numel() for p in net.parameters()), expected_params) + + def test_drop_path_schedule(self): + """Stochastic depth should increase linearly from 0 over the blocks of the whole network.""" + net = ConvNeXt( + spatial_dims=3, + in_channels=1, + out_channels=2, + depths=(1, 1, 2, 1), + features=SMALL_FEATURES, + drop_path_rate=0.3, + ) + rates = [ + getattr(b.drop_path, "drop_prob", 0.0) for i in range(1, 5) for b in getattr(net.features, f"stage{i}") + ] + self.assertEqual(rates, sorted(rates)) + self.assertEqual(rates[0], 0.0) + self.assertAlmostEqual(rates[-1], 0.3) + # a zero rate should not add a drop path module at all + self.assertIsInstance(net.features.stage1[0].drop_path, nn.Identity) + + def test_layer_scale(self): + """`layer_scale_init_value` should control whether the residual branch is scaled.""" + net = ConvNeXt(spatial_dims=2, in_channels=1, out_channels=2, features=SMALL_FEATURES) + self.assertIsNotNone(net.features.stage1[0].gamma) + net = ConvNeXt( + spatial_dims=2, in_channels=1, out_channels=2, features=SMALL_FEATURES, layer_scale_init_value=0.0 + ) + self.assertIsNone(net.features.stage1[0].gamma) + + @parameterized.expand([[2], [3]]) + def test_layer_norm_nd(self, spatial_dims): + """`LayerNormNd` should equal `nn.LayerNorm` applied to the equivalent channels-last tensor.""" + num_channels = 6 + norm_nd = LayerNormNd(num_channels, spatial_dims=spatial_dims) + norm_ref = nn.LayerNorm(num_channels, eps=1e-6) + with torch.no_grad(): + norm_nd.weight.normal_() + norm_nd.bias.normal_() + norm_ref.weight.copy_(norm_nd.weight) + norm_ref.bias.copy_(norm_nd.bias) + + x = torch.randn(2, num_channels, *([4] * spatial_dims)) + channel_last = list(range(spatial_dims + 2)) + channel_last.append(channel_last.pop(1)) # (batch, channel, *spatial) -> (batch, *spatial, channel) + expected = norm_ref(x.permute(channel_last)).permute([0, spatial_dims + 1, *range(1, spatial_dims + 1)]) + assert torch.allclose(norm_nd(x), expected, atol=1e-5) + + def test_ill_arg(self): + with self.assertRaises(ValueError): # unsupported spatial_dims + ConvNeXt(spatial_dims=4, in_channels=1, out_channels=2) + with self.assertRaises(ValueError): # depths and features of different lengths + ConvNeXt(spatial_dims=3, in_channels=1, out_channels=2, depths=(1, 1), features=SMALL_FEATURES) + + +if __name__ == "__main__": + unittest.main() From bb80ddb89822f22871d0ddc1b9307936f933a3c9 Mon Sep 17 00:00:00 2001 From: VenkateswarluNagineni Date: Thu, 16 Jul 2026 08:05:27 -0500 Subject: [PATCH 2/6] feat(networks): validate ConvNeXt arguments and exercise aliases in tests 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 --- monai/networks/nets/convnext.py | 19 ++++++++++++++++--- tests/networks/nets/test_convnext.py | 26 +++++++++++++++++++++++--- 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/monai/networks/nets/convnext.py b/monai/networks/nets/convnext.py index 303b6e0539..634b999842 100644 --- a/monai/networks/nets/convnext.py +++ b/monai/networks/nets/convnext.py @@ -157,7 +157,10 @@ class ConvNeXt(nn.Module): Raises: ValueError: when `spatial_dims` is not one of (1, 2, 3). - ValueError: when `depths` and `features` have different lengths. + ValueError: when `depths` and `features` are empty or of different lengths. + ValueError: when any value in `depths` or `features` is not positive. + ValueError: when `kernel_size` is even, as the block could not preserve its input size. + ValueError: when `drop_path_rate` is outside the interval [0, 1]. Example:: @@ -182,10 +185,20 @@ def __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): + if len(depths) == 0 or len(depths) != len(features): raise ValueError( - f"`depths` and `features` should have the same length, got {len(depths)} and {len(features)}." + f"`depths` and `features` should be non-empty and of the same length, " + f"got lengths {len(depths)} and {len(features)}." ) + if any(d <= 0 for d in depths) or any(f <= 0 for f in features): + raise ValueError(f"`depths` and `features` should be positive, got {tuple(depths)} and {tuple(features)}.") + if kernel_size % 2 == 0: + # an even depthwise kernel cannot keep the spatial size at stride 1, which breaks the residual + raise ValueError( + f"`kernel_size` should be odd so that the block preserves its input size, got {kernel_size}." + ) + if not 0 <= drop_path_rate <= 1: + raise ValueError(f"`drop_path_rate` should be between 0 and 1, got {drop_path_rate}.") conv_type: type[nn.Conv1d | nn.Conv2d | nn.Conv3d] = Conv[Conv.CONV, spatial_dims] avg_pool_type: type[nn.AdaptiveAvgPool1d | nn.AdaptiveAvgPool2d | nn.AdaptiveAvgPool3d] = Pool[ diff --git a/tests/networks/nets/test_convnext.py b/tests/networks/nets/test_convnext.py index eada91b356..e899a67dfe 100644 --- a/tests/networks/nets/test_convnext.py +++ b/tests/networks/nets/test_convnext.py @@ -18,7 +18,18 @@ from parameterized import parameterized from monai.networks import eval_mode -from monai.networks.nets import ConvNeXt, ConvNeXtBase, ConvNeXtLarge, ConvNeXtSmall, ConvNeXtTiny, ConvNeXtXLarge +from monai.networks.nets import ( + ConvNeXt, + Convnext_base, + Convnext_xlarge, + ConvNeXtBase, + ConvNeXtLarge, + ConvNeXtSmall, + ConvNeXtTiny, + ConvNeXtXLarge, + convnext_large, + convnext_tiny, +) from monai.networks.nets.convnext import LayerNormNd from tests.test_utils import skip_if_quick, test_script_save @@ -63,9 +74,10 @@ for case in CASES: TEST_CASES.append([ConvNeXt, *case]) -# every variant, over the 2D and 3D cases, exercising the aliases alongside the canonical names +# every variant, over the 2D and 3D cases, mixing canonical class names with the exported lower/mixed +# case aliases so that a broken package export is caught as well as a broken variant TEST_VARIANT_CASES = [] -for model in [ConvNeXtTiny, ConvNeXtSmall, ConvNeXtBase, ConvNeXtLarge, ConvNeXtXLarge]: +for model in [convnext_tiny, ConvNeXtSmall, Convnext_base, convnext_large, Convnext_xlarge]: for case in [TEST_CASE_1, TEST_CASE_2]: TEST_VARIANT_CASES.append([model, *case]) @@ -149,6 +161,14 @@ def test_ill_arg(self): ConvNeXt(spatial_dims=4, in_channels=1, out_channels=2) with self.assertRaises(ValueError): # depths and features of different lengths ConvNeXt(spatial_dims=3, in_channels=1, out_channels=2, depths=(1, 1), features=SMALL_FEATURES) + with self.assertRaises(ValueError): # empty depths and features + ConvNeXt(spatial_dims=3, in_channels=1, out_channels=2, depths=(), features=()) + with self.assertRaises(ValueError): # non-positive depth + ConvNeXt(spatial_dims=3, in_channels=1, out_channels=2, depths=(1, 0, 1, 1), features=SMALL_FEATURES) + with self.assertRaises(ValueError): # even kernel_size cannot preserve the spatial size + ConvNeXt(spatial_dims=3, in_channels=1, out_channels=2, features=SMALL_FEATURES, kernel_size=4) + with self.assertRaises(ValueError): # drop_path_rate out of range + ConvNeXt(spatial_dims=3, in_channels=1, out_channels=2, features=SMALL_FEATURES, drop_path_rate=1.5) if __name__ == "__main__": From db48c65f794a31d3ed68d68b62cf864e82ce7da8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:05:44 +0000 Subject: [PATCH 3/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/networks/nets/test_convnext.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/networks/nets/test_convnext.py b/tests/networks/nets/test_convnext.py index e899a67dfe..483df917ca 100644 --- a/tests/networks/nets/test_convnext.py +++ b/tests/networks/nets/test_convnext.py @@ -23,10 +23,8 @@ Convnext_base, Convnext_xlarge, ConvNeXtBase, - ConvNeXtLarge, ConvNeXtSmall, ConvNeXtTiny, - ConvNeXtXLarge, convnext_large, convnext_tiny, ) From b167de5031913d3da9ed5f36253358b446cb060d Mon Sep 17 00:00:00 2001 From: VenkateswarluNagineni Date: Tue, 21 Jul 2026 17:47:55 -0500 Subject: [PATCH 4/6] feat(networks): promote LayerNormNd to a public layer 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 --- docs/source/networks.rst | 5 ++ monai/networks/layers/__init__.py | 1 + monai/networks/layers/layer_norm_nd.py | 46 ++++++++++++++++++ monai/networks/nets/convnext.py | 32 +------------ tests/networks/layers/test_layer_norm_nd.py | 52 +++++++++++++++++++++ tests/networks/nets/test_convnext.py | 19 -------- 6 files changed, 105 insertions(+), 50 deletions(-) create mode 100644 monai/networks/layers/layer_norm_nd.py create mode 100644 tests/networks/layers/test_layer_norm_nd.py diff --git a/docs/source/networks.rst b/docs/source/networks.rst index 17f2ddc273..b93d8734f5 100644 --- a/docs/source/networks.rst +++ b/docs/source/networks.rst @@ -326,6 +326,11 @@ Layers .. automodule:: monai.networks.layers.Norm :members: +`LayerNormNd` +~~~~~~~~~~~~~ +.. autoclass:: LayerNormNd + :members: + `Conv` ~~~~~~ .. automodule:: monai.networks.layers.Conv diff --git a/monai/networks/layers/__init__.py b/monai/networks/layers/__init__.py index 48c10270b1..b5855c1169 100644 --- a/monai/networks/layers/__init__.py +++ b/monai/networks/layers/__init__.py @@ -17,6 +17,7 @@ from .factories import Act, Conv, Dropout, LayerFactory, Norm, Pad, Pool, RelPosEmbedding, split_args from .filtering import BilateralFilter, PHLFilter, TrainableBilateralFilter, TrainableJointBilateralFilter from .gmm import GaussianMixtureModel +from .layer_norm_nd import LayerNormNd from .simplelayers import ( LLTM, ApplyFilter, diff --git a/monai/networks/layers/layer_norm_nd.py b/monai/networks/layers/layer_norm_nd.py new file mode 100644 index 0000000000..c3e90185ed --- /dev/null +++ b/monai/networks/layers/layer_norm_nd.py @@ -0,0 +1,46 @@ +# 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 + +import torch +import torch.nn as nn + + +class LayerNormNd(nn.Module): + """ + Layer normalization over the channel dimension of a channels-first tensor. + + `torch.nn.LayerNorm` normalizes over the trailing dimensions, so it expects a channels-last layout + such as (batch, *spatial, channel). Convolutional feature maps in MONAI are channels-first, + (batch, channel, *spatial), and this module normalizes those over the channel dimension only, + for any number of spatial dimensions. + + Args: + num_channels: number of channels of the input, i.e. the size of dimension 1. + spatial_dims: number of spatial dimensions of the input image. + eps: value added to the denominator for numerical stability. + """ + + 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 forward(self, x: torch.Tensor) -> torch.Tensor: + mean = x.mean(dim=1, keepdim=True) + var = (x - mean).pow(2).mean(dim=1, keepdim=True) + x = (x - mean) / torch.sqrt(var + self.eps) + return self.weight.view(self.param_shape) * x + self.bias.view(self.param_shape) diff --git a/monai/networks/nets/convnext.py b/monai/networks/nets/convnext.py index 634b999842..1dfef5f9cc 100644 --- a/monai/networks/nets/convnext.py +++ b/monai/networks/nets/convnext.py @@ -19,6 +19,7 @@ from monai.networks.layers.drop_path import DropPath from monai.networks.layers.factories import Conv, Pool +from monai.networks.layers.layer_norm_nd import LayerNormNd from monai.networks.layers.utils import get_act_layer from monai.networks.layers.weight_init import trunc_normal_ @@ -43,37 +44,6 @@ ] -class LayerNormNd(nn.Module): - """ - Layer normalization over the channel dimension of a channels-first tensor. - - `torch.nn.LayerNorm` normalizes over the trailing dimensions, so it expects a channels-last layout - such as (batch, *spatial, channel). Convolutional feature maps in MONAI are channels-first, - (batch, channel, *spatial), and this module normalizes those over the channel dimension only, - for any number of spatial dimensions. - - Args: - num_channels: number of channels of the input, i.e. the size of dimension 1. - spatial_dims: number of spatial dimensions of the input image. - eps: value added to the denominator for numerical stability. - """ - - 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 forward(self, x: torch.Tensor) -> torch.Tensor: - mean = x.mean(dim=1, keepdim=True) - var = (x - mean).pow(2).mean(dim=1, keepdim=True) - x = (x - mean) / torch.sqrt(var + self.eps) - return self.weight.view(self.param_shape) * x + self.bias.view(self.param_shape) - - class _ConvNeXtBlock(nn.Module): """ A single ConvNeXt block: depthwise convolution, normalization, and an inverted bottleneck of two diff --git a/tests/networks/layers/test_layer_norm_nd.py b/tests/networks/layers/test_layer_norm_nd.py new file mode 100644 index 0000000000..8a6a0aabc5 --- /dev/null +++ b/tests/networks/layers/test_layer_norm_nd.py @@ -0,0 +1,52 @@ +# 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 + +import unittest + +import torch +import torch.nn as nn +from parameterized import parameterized + +from monai.networks.layers import LayerNormNd + + +class TestLayerNormNd(unittest.TestCase): + + @parameterized.expand([[1], [2], [3]]) + def test_shape(self, spatial_dims): + num_channels = 6 + norm = LayerNormNd(num_channels, spatial_dims=spatial_dims) + x = torch.randn(2, num_channels, *([4] * spatial_dims)) + self.assertEqual(norm(x).shape, x.shape) + + @parameterized.expand([[2], [3]]) + def test_equivalent_to_layer_norm(self, spatial_dims): + """`LayerNormNd` should equal `nn.LayerNorm` applied to the equivalent channels-last tensor.""" + num_channels = 6 + norm_nd = LayerNormNd(num_channels, spatial_dims=spatial_dims) + norm_ref = nn.LayerNorm(num_channels, eps=1e-6) + with torch.no_grad(): + norm_nd.weight.normal_() + norm_nd.bias.normal_() + norm_ref.weight.copy_(norm_nd.weight) + norm_ref.bias.copy_(norm_nd.bias) + + x = torch.randn(2, num_channels, *([4] * spatial_dims)) + channel_last = list(range(spatial_dims + 2)) + channel_last.append(channel_last.pop(1)) # (batch, channel, *spatial) -> (batch, *spatial, channel) + expected = norm_ref(x.permute(channel_last)).permute([0, spatial_dims + 1, *range(1, spatial_dims + 1)]) + self.assertTrue(torch.allclose(norm_nd(x), expected, atol=1e-5)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/networks/nets/test_convnext.py b/tests/networks/nets/test_convnext.py index 483df917ca..fe7b9bc740 100644 --- a/tests/networks/nets/test_convnext.py +++ b/tests/networks/nets/test_convnext.py @@ -28,7 +28,6 @@ convnext_large, convnext_tiny, ) -from monai.networks.nets.convnext import LayerNormNd from tests.test_utils import skip_if_quick, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -136,24 +135,6 @@ def test_layer_scale(self): ) self.assertIsNone(net.features.stage1[0].gamma) - @parameterized.expand([[2], [3]]) - def test_layer_norm_nd(self, spatial_dims): - """`LayerNormNd` should equal `nn.LayerNorm` applied to the equivalent channels-last tensor.""" - num_channels = 6 - norm_nd = LayerNormNd(num_channels, spatial_dims=spatial_dims) - norm_ref = nn.LayerNorm(num_channels, eps=1e-6) - with torch.no_grad(): - norm_nd.weight.normal_() - norm_nd.bias.normal_() - norm_ref.weight.copy_(norm_nd.weight) - norm_ref.bias.copy_(norm_nd.bias) - - x = torch.randn(2, num_channels, *([4] * spatial_dims)) - channel_last = list(range(spatial_dims + 2)) - channel_last.append(channel_last.pop(1)) # (batch, channel, *spatial) -> (batch, *spatial, channel) - expected = norm_ref(x.permute(channel_last)).permute([0, spatial_dims + 1, *range(1, spatial_dims + 1)]) - assert torch.allclose(norm_nd(x), expected, atol=1e-5) - def test_ill_arg(self): with self.assertRaises(ValueError): # unsupported spatial_dims ConvNeXt(spatial_dims=4, in_channels=1, out_channels=2) From 4522db264f95483c1a8c4d6286ca5279e4181cfa Mon Sep 17 00:00:00 2001 From: VenkateswarluNagineni Date: Tue, 21 Jul 2026 19:21:02 -0500 Subject: [PATCH 5/6] fix(networks): satisfy mypy and docs build for ConvNeXt 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 --- docs/source/networks.rst | 10 +++++----- monai/networks/nets/convnext.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/networks.rst b/docs/source/networks.rst index b93d8734f5..e5f0aee406 100644 --- a/docs/source/networks.rst +++ b/docs/source/networks.rst @@ -326,11 +326,6 @@ Layers .. automodule:: monai.networks.layers.Norm :members: -`LayerNormNd` -~~~~~~~~~~~~~ -.. autoclass:: LayerNormNd - :members: - `Conv` ~~~~~~ .. automodule:: monai.networks.layers.Conv @@ -348,6 +343,11 @@ Layers .. currentmodule:: monai.networks.layers +`LayerNormNd` +~~~~~~~~~~~~~ +.. autoclass:: LayerNormNd + :members: + `ChannelPad` ~~~~~~~~~~~~ .. autoclass:: ChannelPad diff --git a/monai/networks/nets/convnext.py b/monai/networks/nets/convnext.py index 1dfef5f9cc..86ca514f2d 100644 --- a/monai/networks/nets/convnext.py +++ b/monai/networks/nets/convnext.py @@ -228,7 +228,7 @@ def __init__( if isinstance(m, (conv_type, nn.Linear)): trunc_normal_(m.weight, std=0.02) if m.bias is not None: - nn.init.constant_(m.bias, 0) + nn.init.constant_(torch.as_tensor(m.bias), 0) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.features(x) From 998e5c14576022fe0da153ac41c5e62a78c9e635 Mon Sep 17 00:00:00 2001 From: VenkateswarluNagineni Date: Wed, 22 Jul 2026 07:38:47 -0500 Subject: [PATCH 6/6] fix(networks): render LayerNormNd shape notation as literals 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 --- monai/networks/layers/layer_norm_nd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monai/networks/layers/layer_norm_nd.py b/monai/networks/layers/layer_norm_nd.py index c3e90185ed..97e7bdbbb9 100644 --- a/monai/networks/layers/layer_norm_nd.py +++ b/monai/networks/layers/layer_norm_nd.py @@ -20,8 +20,8 @@ class LayerNormNd(nn.Module): Layer normalization over the channel dimension of a channels-first tensor. `torch.nn.LayerNorm` normalizes over the trailing dimensions, so it expects a channels-last layout - such as (batch, *spatial, channel). Convolutional feature maps in MONAI are channels-first, - (batch, channel, *spatial), and this module normalizes those over the channel dimension only, + such as ``(batch, *spatial, channel)``. Convolutional feature maps in MONAI are channels-first, + ``(batch, channel, *spatial)``, and this module normalizes those over the channel dimension only, for any number of spatial dimensions. Args: