fix(moe): preserve DynamicRoutingLayer compatibility with released checkpoints - #158
Merged
isLinXu merged 1 commit intoJul 22, 2026
Merged
Conversation
There was a problem hiding this comment.
Pull request overview
This PR restores backward compatibility for legacy pickled DynamicRoutingLayer objects found in released checkpoints by avoiding direct reliance on a missing in_channels instance attribute during forward() input validation.
Changes:
- Update
DynamicRoutingLayer.forward()to derive expected input channels via a fallback whenin_channelsis absent. - Add regression tests to ensure legacy objects (missing
in_channels) preserve outputs/state and still enforce channel-shape validation.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| ultralytics/nn/modules/moe/routers.py | Adds compatibility fallback for determining expected router input channels when legacy checkpoints lack in_channels. |
| tests/test_moe_router_boundaries.py | Adds targeted tests covering legacy missing-attribute behavior (output/state invariance and channel-check enforcement). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+444
to
+445
| expected_channels = getattr(self, "in_channels", _get_router_in_channels(self.routing_network)) | ||
| _validate_router_input(x, expected_channels, "DynamicRoutingLayer") |
Collaborator
|
LGTM |
kub-inst
pushed a commit
to kub-inst/YOLO-Master
that referenced
this pull request
Sep 4, 2026
…ter-compat fix(moe): preserve DynamicRoutingLayer compatibility with released checkpoints
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
说明
DynamicRoutingLayer增加兼容处理:当实例中没有in_channels属性时,从routing_network的第一个卷积层读取输入通道数;in_channels,保持当前校验逻辑不变;问题背景
官方 v26.02
YOLO-Master-EsMoE-N.pt中包含四个较早版本序列化的DynamicRoutingLayer。这些对象生成时还没有self.in_channels属性,而当前forward()会直接读取该属性。因此 checkpoint 本身可以加载,但第一次执行 forward 时会报错:所需的输入通道数实际上已经保存在
routing_network的第一个Conv2d中;当前模块也已有_get_router_in_channels(),可以安全读取这一结构信息。修改方式
当
in_channels存在时继续使用原值;只有旧对象缺少该属性时,才回退到_get_router_in_channels(self.routing_network)。该修改不会增加或改变任何 parameter、buffer,也不会改变 checkpoint key。验证结果
2 passed;tests/test_moe_router_boundaries.py:36 passed;tests/test_mixture_fixes.py:10 passed;29e1b93f09b16c8cf7c402f36dcaafc19d4812155631ed45b769e941e4c88c32;AttributeError;64/128/128/256;max_abs_diff均为0.0;影响范围
这是一个针对已发布 pickle 模型对象的向后兼容修复,不会改变 routing weights、Top-K 行为、训练逻辑、导出逻辑或 checkpoint schema。
Related to #52