Block-level ResidualEncoderUNet.yml + clean-key checkpoint (forward-exact 0.0)

body/Body.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:100c7b0f856305031c5e60d55aa3ff46cdb543ec3d52c26e0dea9d9f9753c576
-size 47390413
+oid sha256:ba1c39f84a6128e6813e17f0f193330719fdb91a4d1aa77dce0c920cff78b4d2
+size 47390541

body/Prediction.yml

@@ -1,39 +1,7 @@
 Predictor:
   Model:
-    classpath: segmentation.residualencoderunet.ResidualEncoderUNet
+    classpath: ResidualEncoderUNet.yml
     ResidualEncoderUNet:
-      dim: 2
-      in_channels: 5
-      n_stages: 6
-      features_per_stage:
-      - 24
-      - 48
-      - 96
-      - 192
-      - 256
-      - 256
-      strides:
-      - 1
-      - 2
-      - 2
-      - 2
-      - 2
-      - 2
-      n_blocks_per_stage:
-      - 1
-      - 2
-      - 2
-      - 3
-      - 3
-      - 3
-      n_conv_per_stage_decoder:
-      - 1
-      - 1
-      - 1
-      - 1
-      - 1
-      num_classes: 12
-      deep_supervision: false
       outputs_criterions: None
   Dataset:
     groups_src:

body/ResidualEncoderUNet.yml

@@ -0,0 +1,263 @@
+# Declarative KonfAI ResidualEncoderUNet -- weight-exact with nnU-Net's ResEnc U-Net.
+#
+# This graph reproduces, block-for-block and in forward-execution order, the
+# reference `dynamic_network_architectures.architectures.unet.ResidualEncoderUNet`
+# (the nnU-Net "ResEnc" backbone -- e.g. the ImpactSeg body model). It is the
+# residual-encoder counterpart of `PlainConvUNet.yml`: the DECODER is identical
+# (nnU-Net shares `UNetDecoder` between the two backbones), only the encoder
+# differs. Because the weighted leaves execute in exactly the same sequence as
+# the reference, a REAL ResEnc nnU-Net checkpoint of this topology loads straight
+# into this KonfAI model through the execution-order bridge
+# `konfai.utils.pretrained.transfer_weights_by_execution_order`, and the KonfAI
+# logits are `torch.allclose` with the reference output (see
+# tests/unit/test_residualencoderunet_yaml.py, which transfers the parametric
+# ResidualEncoderUNet's weights in and checks maxdiff < 1e-4).
+#
+# It is authored at the STAGE level from two generic, registry-backed composite
+# blocks -- so the whole ~20-node graph reads as the architecture's topology,
+# not a conv-by-conv unroll:
+#
+#   * `ResidualStage` = one encoder resolution stage: a stack of `n_blocks`
+#     nnU-Net ResNet-D `ResidualBlockD`. Its FIRST block carries the stage stride
+#     (strided-conv downsampling; stage 0 keeps stride 1 = full resolution) and
+#     the channel change; the rest are stride 1. A stride-only stage (channels
+#     unchanged, e.g. 256->256) uses the avgpool residual skip and carries no
+#     projection weights.
+#   * `DecoderStage` = one decoder resolution stage, a TWO-input node
+#     `in_branch: [coarser, skip]`: ConvTranspose(kernel = stride = the matching
+#     encoder stride) on the coarser feature, concat of the encoder skip
+#     (transpose output FIRST, then skip), then `n_conv` decoder conv blocks
+#     mapping 2*skip -> skip. Each conv block is Conv -> InstanceNorm(affine=True)
+#     -> LeakyReLU(0.01), conv_bias=True -- byte-for-byte nnU-Net's signature.
+#
+# The ImpactSeg "body" topology (the default parameters below): a 6-stage ResEnc
+# with feature widths [24, 48, 96, 192, 256, 256], strides [1, 2, 2, 2, 2, 2],
+# n_blocks_per_stage [1, 2, 2, 3, 3, 3], one conv per decoder stage, a 5-channel
+# 2D input and 12 output classes. The single 1x1 seg head emits RAW LOGITS at
+# full resolution (deep_supervision=False -- what nnU-Net returns), which is the
+# weight-exact comparison point.
+#
+# ANISOTROPIC STRIDES/KERNELS: `strides` and `kernel_sizes` accept BOTH isotropic
+# ints AND per-axis lists per stage (e.g. [1, 2, 2]), passed straight through to
+# the convs and the transpose, so real anisotropic ResEnc plans load unchanged.
+#
+# FIXED DEPTH: a declarative YAML cannot loop, so the 6 encoder + 5 decoder stages
+# are structural. To load a checkpoint of a different depth, replicate the stage
+# pattern below (the wiring is identical at every depth); feature widths, strides,
+# block counts, kernels, dim, in_channels and class count stay parametric.
+#
+# Use it as a KonfAI model classpath:
+#   classpath: default|ResidualEncoderUNet.yml   # shipped catalog copy
+#   classpath: /path/to/ResidualEncoderUNet.yml  # any local copy
+# overriding `dim` / `in_channels` / `features_per_stage` / `strides` /
+# `n_blocks_per_stage` / `n_conv_per_stage_decoder` / `num_classes` through the
+# model `parameters`.
+name: ResidualEncoderUNet
+
+parameters:
+  dim: 2
+  in_channels: 5
+  # Encoder feature width per stage; the last entry is the bottleneck. Stage 5
+  # keeps 256->256 (stride-only), which uses the avgpool residual skip.
+  features_per_stage: [24, 48, 96, 192, 256, 256]
+  # Per-stage stride. Entry 0 is the full-resolution stage (stride 1). Each entry
+  # may be an int (isotropic) or a per-axis list (anisotropic, e.g. [1, 2, 2]).
+  strides: [1, 2, 2, 2, 2, 2]
+  # Residual blocks stacked per encoder stage (first carries the stride).
+  n_blocks_per_stage: [1, 2, 2, 3, 3, 3]
+  # Decoder conv blocks per stage (nnU-Net's n_conv_per_stage_decoder).
+  n_conv_per_stage_decoder: [1, 1, 1, 1, 1]
+  num_classes: 12
+  kernel_size: 3
+  conv_bias: true
+  negative_slope: 0.01
+
+network:
+  in_channels: ${in_channels}
+  dim: ${dim}
+
+modules:
+  # ----- Stem: one non-strided conv block (never downsamples), 5 -> 24 ---------
+  - name: Stem
+    type: ConvBlock
+    in_branch: [0]
+    out_branch: [stem]
+    args:
+      dim: ${dim}
+      in_channels: ${in_channels}
+      out_channels: ${features_per_stage.0}
+      block_configs:
+        - $object: BlockConfig
+          args:
+            kernel_size: ${kernel_size}
+            stride: 1
+            padding: 1
+            bias: ${conv_bias}
+            activation: LeakyReLU;0.01
+            norm_mode: INSTANCE_AFFINE
+
+  # ----- Encoder: residual stages (each output is a decoder skip) --------------
+  # Stage k: stack n_blocks_per_stage[k] ResidualBlockD; first block strides + changes width.
+  - name: Encoder_0
+    type: ResidualStage
+    in_branch: [stem]
+    out_branch: [enc0]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.0}
+      out_channels: ${features_per_stage.0}
+      n_blocks: ${n_blocks_per_stage.0}
+      stride: ${strides.0}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+  - name: Encoder_1
+    type: ResidualStage
+    in_branch: [enc0]
+    out_branch: [enc1]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.0}
+      out_channels: ${features_per_stage.1}
+      n_blocks: ${n_blocks_per_stage.1}
+      stride: ${strides.1}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+  - name: Encoder_2
+    type: ResidualStage
+    in_branch: [enc1]
+    out_branch: [enc2]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.1}
+      out_channels: ${features_per_stage.2}
+      n_blocks: ${n_blocks_per_stage.2}
+      stride: ${strides.2}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+  - name: Encoder_3
+    type: ResidualStage
+    in_branch: [enc2]
+    out_branch: [enc3]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.2}
+      out_channels: ${features_per_stage.3}
+      n_blocks: ${n_blocks_per_stage.3}
+      stride: ${strides.3}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+  - name: Encoder_4
+    type: ResidualStage
+    in_branch: [enc3]
+    out_branch: [enc4]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.3}
+      out_channels: ${features_per_stage.4}
+      n_blocks: ${n_blocks_per_stage.4}
+      stride: ${strides.4}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+  # Stage 5: bottleneck, 256 -> 256 at stride 2 (avgpool skip, no projection).
+  - name: Encoder_5
+    type: ResidualStage
+    in_branch: [enc4]
+    out_branch: [enc5]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.4}
+      out_channels: ${features_per_stage.5}
+      n_blocks: ${n_blocks_per_stage.5}
+      stride: ${strides.5}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+
+  # ----- Decoder: coarsest-first, upsample + concat skip + conv blocks ---------
+  # Stage j upsamples encoder feature (n_stages-1-j) and concatenates skip (n_stages-2-j).
+  - name: Decoder_0
+    type: DecoderStage
+    in_branch: [enc5, enc4]
+    out_branch: [dec0]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.5}
+      skip_channels: ${features_per_stage.4}
+      n_conv: ${n_conv_per_stage_decoder.0}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+      upsample_stride: ${strides.5}
+  - name: Decoder_1
+    type: DecoderStage
+    in_branch: [dec0, enc3]
+    out_branch: [dec1]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.4}
+      skip_channels: ${features_per_stage.3}
+      n_conv: ${n_conv_per_stage_decoder.1}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+      upsample_stride: ${strides.4}
+  - name: Decoder_2
+    type: DecoderStage
+    in_branch: [dec1, enc2]
+    out_branch: [dec2]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.3}
+      skip_channels: ${features_per_stage.2}
+      n_conv: ${n_conv_per_stage_decoder.2}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+      upsample_stride: ${strides.3}
+  - name: Decoder_3
+    type: DecoderStage
+    in_branch: [dec2, enc1]
+    out_branch: [dec3]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.2}
+      skip_channels: ${features_per_stage.1}
+      n_conv: ${n_conv_per_stage_decoder.3}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+      upsample_stride: ${strides.2}
+  - name: Decoder_4
+    type: DecoderStage
+    in_branch: [dec3, enc0]
+    out_branch: [dec4]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.1}
+      skip_channels: ${features_per_stage.0}
+      n_conv: ${n_conv_per_stage_decoder.4}
+      kernel_size: ${kernel_size}
+      conv_bias: ${conv_bias}
+      negative_slope: ${negative_slope}
+      upsample_stride: ${strides.1}
+
+  # ----- Segmentation head -----------------------------------------------------
+  # Single full-resolution 1x1 conv emitting RAW LOGITS (deep_supervision=False,
+  # what nnU-Net returns). This is the weight-exact comparison point.
+  - name: SegHead
+    type: Conv
+    in_branch: [dec4]
+    out_branch: [-1]
+    args:
+      dim: ${dim}
+      in_channels: ${features_per_stage.0}
+      out_channels: ${num_classes}
+      kernel_size: 1
+      stride: 1
+      padding: 0
+      bias: true