fix(moe): preserve pruned experts' kernel sizes in model.yaml so retraining reloads their weights

tests/test_moe_prune_yaml_sync.py

@@ -1,14 +1,18 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
-"""Regression test: MoEPruner must keep ``model.yaml`` consistent with the pruned
-expert count so a pruned model survives a YAML-based rebuild during retraining
+"""Regression tests: MoEPruner must keep ``model.yaml`` consistent with the pruned
+ES_MOE blocks so a pruned model survives a YAML-based rebuild during retraining
 (the prune -> LoRA / full fine-tune recovery workflow).
 
-Without the fix, ``DetectionModel(pruned.yaml)`` rebuilds with ES_MOE's default
-expert count and ``intersect_dicts`` silently drops the reduced expert/router
-weights, re-inflating the model with randomly initialized experts.
-"""
-import copy
+Two failure modes are covered:
+1. Expert *count* re-inflation - rebuild uses ES_MOE's default expert count and
+   ``intersect_dicts`` drops the reduced expert/router weights.
+2. Expert *kernel* mismatch - pruning keeps experts with heterogeneous kernels
+   (e.g. [5, 9]) but a bare rebuild assigns the defaults [3, 5], so the kept
+   experts' depthwise weights are dropped on a shape mismatch and re-initialized.
 
+The fix writes the full ES_MOE arg list (count + per-expert kernel sizes) into
+``model.yaml`` so both survive the rebuild.
+"""
 import yaml
 
 from ultralytics.nn.modules.moe.modules import ES_MOE
@@ -22,6 +26,14 @@ def _experts(model):
     return [m.num_experts for _, m in model.named_modules() if isinstance(m, ES_MOE)]
 
 
+def _kernels(model):
+    return [
+        [e.conv.depthwise.kernel_size[0] for e in m.experts]
+        for _, m in model.named_modules()
+        if isinstance(m, ES_MOE)
+    ]
+
+
 def _build_reduced(num_experts):
     """Build a model whose ES_MOE blocks each have ``num_experts`` (stand-in for a pruned model)."""
     d = yaml.safe_load(open(CFG))
@@ -31,6 +43,16 @@ def _build_reduced(num_experts):
     return DetectionModel(d, ch=3, nc=10, verbose=False)
 
 
+def _build_with_kernels(num_experts, kernels):
+    """Stand-in for a pruned model that kept experts with non-default kernel sizes."""
+    d = yaml.safe_load(open(CFG))
+    for layer in d["backbone"]:
+        if layer[2] == "ES_MOE":
+            out_ch = layer[3][0]
+            layer[3] = [out_ch, num_experts, 8, num_experts, True, 0.4, 15, list(kernels)]
+    return DetectionModel(d, ch=3, nc=10, verbose=False)
+
+
 def test_reinflation_without_sync():
     """Reproduce the bug: reduced model + un-synced (full) yaml re-inflates on rebuild."""
     reduced = _build_reduced(2)
@@ -41,12 +63,24 @@ def test_reinflation_without_sync():
 
 
 def test_sync_preserves_pruned_experts():
-    """The fix: MoEPruner._sync_yaml_num_experts writes the reduced count into yaml,
-    so a YAML rebuild preserves the pruned architecture."""
+    """The fix writes the reduced count into yaml so a rebuild preserves it."""
     reduced = _build_reduced(2)
     reduced.yaml = yaml.safe_load(open(CFG))  # start from the un-synced (buggy) yaml
     MoEPruner._sync_yaml_num_experts(type("D", (), {})(), reduced)  # method only uses self for logging
     es_args = [layer[3] for layer in reduced.yaml["backbone"] if layer[2] == "ES_MOE"]
-    assert all(a[-1] == 2 for a in es_args), f"yaml not synced: {es_args}"
+    assert all(a[1] == 2 for a in es_args), f"num_experts not synced: {es_args}"
     rebuilt = DetectionModel(reduced.yaml, ch=3, nc=10, verbose=False)
     assert _experts(rebuilt) == [2, 2, 2, 2], "pruned expert count must survive the rebuild"
+
+
+def test_sync_preserves_expert_kernels():
+    """The fix also writes per-expert kernel sizes so heterogeneous kept experts survive."""
+    reduced = _build_with_kernels(2, [5, 9])
+    before = _kernels(reduced)
+    assert all(k == [5, 9] for k in before), f"stand-in kernels not built: {before}"
+    reduced.yaml = yaml.safe_load(open(CFG))  # un-synced yaml would rebuild default [3, 5]
+    MoEPruner._sync_yaml_num_experts(type("D", (), {})(), reduced)
+    es_args = [layer[3] for layer in reduced.yaml["backbone"] if layer[2] == "ES_MOE"]
+    assert all(a[-1] == [5, 9] for a in es_args), f"kernels not synced: {es_args}"
+    rebuilt = DetectionModel(reduced.yaml, ch=3, nc=10, verbose=False)
+    assert _kernels(rebuilt) == before, "expert kernel sizes must survive the rebuild"

ultralytics/nn/modules/moe/modules.py

@@ -404,7 +404,7 @@ class ES_MOE(nn.Module):
 
     def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8,
                  top_k=2, use_sparse_inference=True, dynamic_threshold=0.4,
-                 max_kernel_size=15):
+                 max_kernel_size=15, expert_kernel_sizes=None):
         """
         Args:
             in_channels: Input channels
@@ -415,6 +415,10 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8,
             use_sparse_inference: Enable sparse Top-K expert computation during inference
             dynamic_threshold: Optional threshold for pruning low-confidence experts during inference
             max_kernel_size: Largest odd depthwise kernel assigned to an expert
+            expert_kernel_sizes: Optional explicit per-expert depthwise kernel sizes
+                (length must equal ``num_experts``). When ``None`` the kernels are
+                derived from defaults; pruned checkpoints set this so retraining
+                rebuilds the exact kept-expert kernels and reloads their weights.
         """
         super(ES_MOE, self).__init__()
 
@@ -440,6 +444,7 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8,
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.num_experts = num_experts
+        self.reduction = reduction
         self.top_k = min(top_k, num_experts) if top_k is not None else num_experts
         self.use_top_k = (top_k is not None)
         self.use_sparse_inference = use_sparse_inference
@@ -449,12 +454,27 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8,
         # Dynamic routing (Top-K supported)
         self.routing = DynamicRoutingLayer(in_channels, num_experts, reduction, top_k)
 
-        # Expert group (original design)
-        default_kernel_sizes = [3, 5, 7]
-        if num_experts <= len(default_kernel_sizes):
-            ks = [min(k, max_kernel_size) for k in default_kernel_sizes[:num_experts]]
+        # Expert group (original design). ``expert_kernel_sizes`` lets a pruned
+        # checkpoint reconstruct its kept experts' heterogeneous kernels so that
+        # ``YOLO(pruned.pt).train()`` reloads expert weights instead of dropping
+        # them on a kernel-shape mismatch (prune -> LoRA/full fine-tune recovery).
+        if expert_kernel_sizes is not None:
+            if len(expert_kernel_sizes) != num_experts:
+                raise ValueError(
+                    f"expert_kernel_sizes must have {num_experts} entries, got {len(expert_kernel_sizes)}"
+                )
+            ks = []
+            for k in expert_kernel_sizes:
+                k = int(k)
+                if k % 2 == 0:
+                    k -= 1
+                ks.append(min(k, max_kernel_size))
         else:
-            ks = [min(3 + 2 * i, max_kernel_size) for i in range(num_experts)]
+            default_kernel_sizes = [3, 5, 7]
+            if num_experts <= len(default_kernel_sizes):
+                ks = [min(k, max_kernel_size) for k in default_kernel_sizes[:num_experts]]
+            else:
+                ks = [min(3 + 2 * i, max_kernel_size) for i in range(num_experts)]
         self.experts = nn.ModuleList(
             [EfficientExpertGroup(in_channels, out_channels, kernel_size=k) for k in ks]
         )

ultralytics/nn/modules/moe/pruning.py

@@ -442,8 +442,30 @@ def _sync_yaml_num_experts(self, pruned_model: nn.Module) -> None:
                 continue
             args = list(seq[j][3]) if isinstance(seq[j][3], list) else [seq[j][3]]
             out_ch = args[0] if args else getattr(mod, "out_channels", None)
-            seq[j][3] = [out_ch, int(mod.num_experts)]
-        LOGGER.info("   Synced post-prune expert counts into model.yaml")
+            # Recover the kept experts' actual depthwise kernel sizes; pruning keeps
+            # experts with heterogeneous kernels whose order/values are not the
+            # default [3, 5, 7...] a bare rebuild would assign. Writing the full
+            # positional arg list (incl. expert_kernel_sizes) lets YOLO(pruned.pt)
+            # .train() rebuild the exact experts so their weights survive
+            # intersect_dicts instead of being dropped on a kernel-shape mismatch.
+            kernels = []
+            for expert in mod.experts:
+                conv = getattr(expert, "conv", None)
+                depthwise = getattr(conv, "depthwise", None) if conv is not None else None
+                ks = getattr(depthwise, "kernel_size", None)
+                kernels.append(int(ks[0]) if isinstance(ks, (tuple, list)) else int(ks) if ks else 3)
+            top_k = int(mod.top_k) if getattr(mod, "use_top_k", True) else None
+            seq[j][3] = [
+                out_ch,
+                int(mod.num_experts),
+                int(getattr(mod, "reduction", 8)),
+                top_k,
+                bool(getattr(mod, "use_sparse_inference", True)),
+                float(getattr(mod, "dynamic_threshold", 0.4)),
+                int(getattr(mod, "max_kernel_size", 15)),
+                kernels,
+            ]
+        LOGGER.info("   Synced post-prune expert counts + kernel sizes into model.yaml")
 
     def _save_model(self, pruned_model: nn.Module, output_path: str) -> None:
         """