Skip to content

ModelnnUNetWrapper fails with nnunetv2 >= 2.7: build_network_architecture signature changed #9239

Open

Description

@garciadias

Describe the bug
ModelnnUNetWrapper.__init__ (and therefore get_nnunet_monai_predictor) calls

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

(monai/apps/nnunet/nnunet_bundle.py, around line 219 on dev). That matches the signature in nnunetv2 2.4–2.6. In nnunetv2 2.7.0 the static method became

build_network_architecture(plans_manager, configuration_manager, num_input_channels, num_output_channels, enable_deep_supervision=True)

so with nnunetv2 2.7.0 or later the positional arguments are shifted and the call fails:

TypeError: nnUNetTrainer.build_network_architecture() got multiple values for argument 'enable_deep_supervision'

get_network_from_nnunet_plans is not affected, because it calls nnunetv2.utilities.get_network_from_plans.get_network_from_plans, whose signature didn't change.

To Reproduce
With nnunetv2==2.8.1 and torch (CPU is enough) installed:

  1. Create a model folder <root>/fold_0 with <root>/plans.json (any standard 2d PlainConvUNet plans), <root>/dataset.json, and <root>/nnunet_checkpoint.pth containing {"trainer_name": "nnUNetTrainer", "init_args": {"configuration": "2d"}}.
  2. Run ModelnnUNetWrapper(predictor, "<root>/fold_0") (or get_nnunet_monai_predictor("<root>/fold_0")).
  3. The TypeError above is raised.

Expected behavior
The wrapper builds the network with any nnunetv2 release that MONAI supports.

Environment
MONAI dev @ 2c54111, nnunetv2 2.8.1 (also checked: the new signature is in 2.7.0 and 2.8.0; 2.4.2–2.6.2 use the old one), torch 2.15.0+cpu, Python 3.12.

Additional context

  • No CI job installs nnunetv2, so tests/integration/test_integration_nnunet_bundle.py doesn't catch this.
  • A possible fix is to choose the call by signature, e.g. pass plans_manager, configuration_manager, num_input_channels, num_output_channels when the first parameter of build_network_architecture is plans_manager, otherwise use the current arguments. Alternatively, pin nnunetv2<2.7 for the bundle helpers.
  • Check nnUNet plans network names against allowed modules #9237 edits the lines just before this call, so whichever merges second may need a small rebase.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions