Add MedNeXt head hidden channels and update tutorials

Author: Donglai Wei

connectomics/config/pipeline/config_io.py

@@ -304,6 +304,9 @@ def validate_config(cfg: Config) -> None:
         for head_name, head_cfg in model_heads.items():
             head_out_channels = int(getattr(head_cfg, "out_channels", 0))
             head_num_blocks = int(getattr(head_cfg, "num_blocks", 0))
+            head_hidden_channels = getattr(head_cfg, "hidden_channels", None)
+            if head_hidden_channels is not None:
+                head_hidden_channels = int(head_hidden_channels)
             if head_out_channels <= 0:
                 raise ValueError(
                     f"model.heads.{head_name}.out_channels must be positive "
@@ -314,6 +317,11 @@ def validate_config(cfg: Config) -> None:
                     f"model.heads.{head_name}.num_blocks must be non-negative "
                     f"(got {head_num_blocks})"
                 )
+            if head_hidden_channels is not None and head_hidden_channels <= 0:
+                raise ValueError(
+                    f"model.heads.{head_name}.hidden_channels must be positive "
+                    f"(got {head_hidden_channels})"
+                )
 
         primary_head = getattr(cfg.model, "primary_head", None)
         if primary_head is not None and primary_head not in model_heads:
@@ -326,7 +334,11 @@ def validate_config(cfg: Config) -> None:
                 f"inference.head='{inference_head}' is not present in model.heads "
                 f"({sorted(model_heads.keys())})."
             )
-        if visualization_head is not None and visualization_head != "all" and visualization_head not in model_heads:
+        if (
+            visualization_head is not None
+            and visualization_head != "all"
+            and visualization_head not in model_heads
+        ):
             raise ValueError(
                 f"monitor.logging.images.head='{visualization_head}' is not present in "
                 f"model.heads ({sorted(model_heads.keys())})."

connectomics/models/architectures/mednext_models.py

@@ -24,8 +24,7 @@
 
 try:
     from nnunet_mednext import MedNeXt as MedNeXtBase
-    from nnunet_mednext import MedNeXtBlock
-    from nnunet_mednext import create_mednext_v1
+    from nnunet_mednext import MedNeXtBlock, create_mednext_v1
 
     MEDNEXT_AVAILABLE = True
 except ImportError:
@@ -137,6 +136,7 @@ def __init__(
         in_channels: int,
         out_channels: int,
         num_blocks: int,
+        hidden_channels: int | None = None,
         *,
         exp_r: int,
         kernel_size: int,
@@ -150,17 +150,33 @@ def __init__(
             raise ValueError(f"MedNeXt task head num_blocks must be >= 0, got {num_blocks}")
         if out_channels <= 0:
             raise ValueError(f"MedNeXt task head out_channels must be positive, got {out_channels}")
+        if hidden_channels is None:
+            hidden_channels = in_channels
+        if hidden_channels <= 0:
+            raise ValueError(
+                f"MedNeXt task head hidden_channels must be positive, got {hidden_channels}"
+            )
+        if hidden_channels > in_channels:
+            raise ValueError(
+                "MedNeXt task head hidden_channels must not exceed the shared feature width "
+                f"({hidden_channels} > {in_channels})"
+            )
         if dim == "2d":
             conv = nn.Conv2d
         elif dim == "3d":
             conv = nn.Conv3d
         else:
             raise ValueError(f"MedNeXt task head dim must be '2d' or '3d', got {dim}")
 
+        self.input_projection = (
+            conv(in_channels, hidden_channels, kernel_size=1)
+            if hidden_channels != in_channels
+            else nn.Identity()
+        )
         blocks = [
             MedNeXtBlock(
-                in_channels=in_channels,
-                out_channels=in_channels,
+                in_channels=hidden_channels,
+                out_channels=hidden_channels,
                 exp_r=exp_r,
                 kernel_size=kernel_size,
                 do_res=do_res,
@@ -171,9 +187,11 @@ def __init__(
             for _ in range(num_blocks)
         ]
         self.blocks = nn.Sequential(*blocks) if blocks else nn.Identity()
-        self.projection = conv(in_channels, out_channels, kernel_size=1)
+        self.projection = conv(hidden_channels, out_channels, kernel_size=1)
+        self.hidden_channels = hidden_channels
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.input_projection(x)
         x = self.blocks(x)
         return self.projection(x)
 
@@ -219,15 +237,19 @@ def __init__(
         for head_name, head_cfg in heads.items():
             out_channels = int(_cfg_value(head_cfg, "out_channels", head_cfg))
             num_blocks = int(_cfg_value(head_cfg, "num_blocks", 0))
+            hidden_channels = _cfg_value(head_cfg, "hidden_channels", None)
+            hidden_channels = int(hidden_channels) if hidden_channels is not None else None
             task_heads[head_name] = MedNeXtTaskHead(
                 in_channels=self.feature_channels,
                 out_channels=out_channels,
                 num_blocks=num_blocks,
+                hidden_channels=hidden_channels,
                 **self.head_block_kwargs,
             )
             head_specs[head_name] = {
                 "out_channels": out_channels,
                 "num_blocks": num_blocks,
+                "hidden_channels": hidden_channels or self.feature_channels,
             }
 
         self.heads = nn.ModuleDict(task_heads)

tests/unit/test_mednext_multi_head_wrapper.py

@@ -1,10 +1,9 @@
 import pytest
 import torch
 
-from connectomics.config import from_dict
-from connectomics.models.build import build_model
+from connectomics.config import from_dict, validate_config
 from connectomics.models.architectures.mednext_models import MedNeXtMultiHeadWrapper
-
+from connectomics.models.build import build_model
 
 nnunet_mednext = pytest.importorskip("nnunet_mednext")
 from nnunet_mednext import MedNeXt  # noqa: E402
@@ -32,7 +31,7 @@ def test_mednext_multi_head_wrapper_returns_named_outputs():
     model = MedNeXtMultiHeadWrapper(
         trunk,
         {
-            "affinity": {"out_channels": 9, "num_blocks": 1},
+            "affinity": {"out_channels": 9, "num_blocks": 1, "hidden_channels": 2},
             "sdt": {"out_channels": 1, "num_blocks": 0},
         },
     )
@@ -51,7 +50,21 @@ def test_mednext_multi_head_wrapper_returns_named_outputs():
     assert outputs["output"]["affinity"].shape == (1, 9, 32, 32, 32)
     assert outputs["output"]["sdt"].shape == (1, 1, 32, 32, 32)
     assert model.head_specs["affinity"]["num_blocks"] == 1
+    assert model.head_specs["affinity"]["hidden_channels"] == 2
+    assert isinstance(model.heads["affinity"].input_projection, torch.nn.Conv3d)
+    assert model.heads["affinity"].projection.in_channels == 2
     assert model.head_specs["sdt"]["num_blocks"] == 0
+    assert model.head_specs["sdt"]["hidden_channels"] == 4
+
+
+def test_mednext_multi_head_wrapper_rejects_hidden_channels_above_trunk_width():
+    trunk = _build_tiny_mednext(deep_supervision=False)
+
+    with pytest.raises(ValueError, match="must not exceed the shared feature width"):
+        MedNeXtMultiHeadWrapper(
+            trunk,
+            {"affinity": {"out_channels": 9, "num_blocks": 1, "hidden_channels": 8}},
+        )
 
 
 def test_mednext_multi_head_wrapper_rejects_deep_supervision_trunk():
@@ -70,7 +83,7 @@ def test_build_model_creates_mednext_multi_head_wrapper_from_config():
                 "out_channels": 10,
                 "primary_head": "affinity",
                 "heads": {
-                    "affinity": {"out_channels": 9, "num_blocks": 1},
+                    "affinity": {"out_channels": 9, "num_blocks": 1, "hidden_channels": 2},
                     "sdt": {"out_channels": 1, "num_blocks": 0},
                 },
                 "mednext": {
@@ -91,10 +104,25 @@ def test_build_model_creates_mednext_multi_head_wrapper_from_config():
     model = build_model(cfg)
     assert isinstance(model, MedNeXtMultiHeadWrapper)
     assert model.primary_head == "affinity"
+    assert model.head_specs["affinity"]["hidden_channels"] == 2
 
     x = torch.randn(1, 1, 32, 32, 32)
     with torch.no_grad():
         outputs = model(x)
 
     assert outputs["output"]["affinity"].shape == (1, 9, 32, 32, 32)
     assert outputs["output"]["sdt"].shape == (1, 1, 32, 32, 32)
+
+
+def test_validate_config_rejects_nonpositive_head_hidden_channels():
+    cfg = from_dict(
+        {
+            "model": {
+                "arch": {"type": "mednext_custom"},
+                "heads": {"affinity": {"out_channels": 3, "num_blocks": 0, "hidden_channels": 0}},
+            }
+        }
+    )
+
+    with pytest.raises(ValueError, match="model.heads.affinity.hidden_channels must be positive"):
+        validate_config(cfg)

tutorials/mito_betaseg.yaml

@@ -18,14 +18,17 @@ default:
       aff_r1:
         out_channels: 3
         num_blocks: 1
+        hidden_channels: 8
         target_slice: "0:3"
       aff_r5:
         out_channels: 3
         num_blocks: 1
+        hidden_channels: 8
         target_slice: "3:6"
       sdt:
         out_channels: 1
         num_blocks: 1
+        hidden_channels: 8
         target_slice: "6:7"
     loss:
       losses:
@@ -60,12 +63,10 @@ default:
             bg_value: -1.0
     dataloader:
       patch_size: [128, 128, 128]
-      batch_size: 4
+      batch_size: 2
       use_lazy_zarr: true
     image_transform:
-      normalize: "0-1"
-      clip_percentile_low: 0.005
-      clip_percentile_high: 0.995
+      normalize: "divide-255"
     augmentation:
       profile: aug_em_neuron
       rotate:
@@ -123,7 +124,7 @@ train:
         channel_mode: all
     checkpoint:
       save_top_k: 3
-      monitor: val_loss_total_epoch
+      monitor: val_loss_total
       mode: min
 
 test:

tutorials/neuron_snemi_sdt_multitask.yaml

@@ -148,7 +148,7 @@ test:
       label: test-labels.h5
       resolution: [30, 6, 6]
   inference:
-    head: aff_r5
+    head: aff_r1
     decoding:
       - template: decoding_waterz
         kwargs:
@@ -160,6 +160,11 @@ test:
           dust_merge_size: 800
           dust_merge_affinity: 0.3
           dust_remove_size: 600
+          branch_merge: true
+          branch_iou_threshold: 0.5
+          branch_best_buddy: true
+          branch_one_sided_threshold: 0.8
+          branch_one_sided_min_size: 100
 
 # ============================================================================
 # Parameter tuning for waterz agglomeration thresholds (--mode tune)