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..427241e872 100644 --- a/tests/apps/nnunet/test_nnunet_bundle_network_spec.py +++ b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py @@ -216,6 +216,64 @@ 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() + + # 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, 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.""" from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer