From 025dd3bbd37fafd7e1bca86a847ae30616187b73 Mon Sep 17 00:00:00 2001 From: shiva Date: Sat, 31 Oct 2026 12:25:21 +0530 Subject: [PATCH 1/2] Support nnunetv2 >= 2.7 build_network_architecture signature nnunetv2 2.7 changed nnUNetTrainer.build_network_architecture to take the plans and configuration managers instead of the unpacked network arguments. ModelnnUNetWrapper still passed the old arguments by position, so loading a model with nnunetv2 2.7 or later raised "TypeError: got multiple values for argument 'enable_deep_supervision'". Choose the call from the trainer's signature, as nnunetv2 does in its own predictor and trainer, so both old and new releases are supported. Fixes #9239. Signed-off-by: shiva --- monai/apps/nnunet/nnunet_bundle.py | 28 +++++++---- .../nnunet/test_nnunet_bundle_network_spec.py | 48 +++++++++++++++++++ 2 files changed, 68 insertions(+), 8 deletions(-) diff --git a/monai/apps/nnunet/nnunet_bundle.py b/monai/apps/nnunet/nnunet_bundle.py index d71a2aa38f..60d0491ca7 100644 --- a/monai/apps/nnunet/nnunet_bundle.py +++ b/monai/apps/nnunet/nnunet_bundle.py @@ -10,6 +10,7 @@ # limitations under the License. from __future__ import annotations +import inspect import os import shutil from collections.abc import Sequence @@ -285,14 +286,25 @@ def __init__( configuration_manager.network_arch_init_kwargs_req_import, allowed_network_modules, ) - network = trainer_class.build_network_architecture( - configuration_manager.network_arch_class_name, - configuration_manager.network_arch_init_kwargs, - configuration_manager.network_arch_init_kwargs_req_import, - num_input_channels, - plans_manager.get_label_manager(dataset_json).num_segmentation_heads, - enable_deep_supervision=False, - ) + num_output_channels = plans_manager.get_label_manager(dataset_json).num_segmentation_heads + # nnunetv2 >= 2.7 trainers take the plans and configuration managers, not the unpacked architecture arguments + if "plans_manager" in inspect.signature(trainer_class.build_network_architecture).parameters: + network = trainer_class.build_network_architecture( + plans_manager, + configuration_manager, + num_input_channels, + num_output_channels, + enable_deep_supervision=False, + ) + else: + network = trainer_class.build_network_architecture( + configuration_manager.network_arch_class_name, + configuration_manager.network_arch_init_kwargs, + configuration_manager.network_arch_init_kwargs_req_import, + num_input_channels, + num_output_channels, + enable_deep_supervision=False, + ) predictor.plans_manager = plans_manager # type: ignore predictor.configuration_manager = configuration_manager # type: ignore diff --git a/tests/apps/nnunet/test_nnunet_bundle_network_spec.py b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py index de0cb2eab2..a4124edfa3 100644 --- a/tests/apps/nnunet/test_nnunet_bundle_network_spec.py +++ b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py @@ -216,6 +216,54 @@ def test_default_plans_pass_the_check(self): self._wrap(_plans()) build.assert_called_once() + def test_default_plans_build_the_network(self): + """The installed nnunetv2 trainer builds the plans' network without deep supervision.""" + network = self._wrap(_plans()).network_weights + self.assertEqual(type(network).__name__, "PlainConvUNet") + self.assertFalse(network.decoder.deep_supervision) + + def test_build_network_architecture_signatures(self): + """Trainers with the nnunetv2 < 2.7 and >= 2.7 ``build_network_architecture`` signatures are both supported.""" + import torch + from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer + + calls = {} + + def legacy( + architecture_class_name, + arch_init_kwargs, + arch_init_kwargs_req_import, + num_input_channels, + num_output_channels, + enable_deep_supervision=True, + ): + calls["legacy"] = ( + architecture_class_name, + num_input_channels, + num_output_channels, + enable_deep_supervision, + ) + return torch.nn.Identity() + + def current( + plans_manager, configuration_manager, num_input_channels, num_output_channels, enable_deep_supervision=True + ): + calls["current"] = ( + configuration_manager.network_arch_class_name, + num_input_channels, + num_output_channels, + enable_deep_supervision, + ) + return torch.nn.Identity() + + for name, build in (("legacy", legacy), ("current", current)): + with ( + self.subTest(name), + mock.patch.object(nnUNetTrainer, "build_network_architecture", staticmethod(build)), + ): + self._wrap(_plans()) + self.assertEqual(calls.get(name), (UNET, 1, 2, False)) + def test_disallowed_network_is_not_built(self): """A network class outside the allowed modules is rejected before the trainer builds anything.""" from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer From 6e3c601b2c50747eef7a8a816581fea36c0a6dab Mon Sep 17 00:00:00 2001 From: shiva Date: Sat, 31 Oct 2026 12:40:30 +0530 Subject: [PATCH 2/2] Check computed channel counts in the nnUNet wrapper signature test Use a dataset with two input channels and three labels so the test fails if the wrapper passes fixed channel counts to the trainer. Signed-off-by: shiva --- tests/apps/nnunet/test_nnunet_bundle_network_spec.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/apps/nnunet/test_nnunet_bundle_network_spec.py b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py index a4124edfa3..427241e872 100644 --- a/tests/apps/nnunet/test_nnunet_bundle_network_spec.py +++ b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py @@ -256,13 +256,23 @@ def current( ) return torch.nn.Identity() + # two input channels and three labels, so the channel counts differ from the single-channel defaults + with open(os.path.join(self._tmp.name, "dataset.json"), "w") as f: + json.dump( + { + "channel_names": {"0": "T1", "1": "T2"}, + "labels": {"background": 0, "a": 1, "b": 2}, + "file_ending": ".nii.gz", + }, + f, + ) for name, build in (("legacy", legacy), ("current", current)): with ( self.subTest(name), mock.patch.object(nnUNetTrainer, "build_network_architecture", staticmethod(build)), ): self._wrap(_plans()) - self.assertEqual(calls.get(name), (UNET, 1, 2, False)) + self.assertEqual(calls.get(name), (UNET, 2, 3, False)) def test_disallowed_network_is_not_built(self): """A network class outside the allowed modules is rejected before the trainer builds anything."""