fix_moe_architecture.py

#!/usr/bin/env python3
"""
修复MoE架构实现问题

Author: Augment Agent (Claude Sonnet 4 by Anthropic)
Created: 2025-07-05
Description: 根据GUIDE/MOE.md文档要求，修复当前MoE实现中的问题
"""

import yaml
from pathlib import Path


def analyze_moe_issues():
    """分析当前MoE实现的问题"""
    print("🔍 分析当前MoE架构实现问题")
    print("=" * 60)
    
    issues = []
    
    # 问题1: YAML配置参数顺序错误
    issues.append({
        'issue': 'YAML配置参数顺序错误',
        'description': 'C3MoE的参数顺序与实际__init__方法不匹配',
        'current': '[-1, 3, C3MoE, [128, 1, True, 1, 0.5, 4, 2]]',
        'correct': '[-1, 3, C3MoE, [128, True, 1, 0.5, 4, 2]]',
        'explanation': '多了一个参数，导致参数错位'
    })
    
    # 问题2: 违背MoE核心原则
    issues.append({
        'issue': '违背MoE核心原则',
        'description': '过度使用MoE层，违背了"在关键位置使用MoE"的原则',
        'current': '几乎所有层都使用MoE',
        'correct': '只在C3模块中使用MoE，保留标准Conv层',
        'explanation': 'GUIDE/MOE.md明确指出应该在"最消耗能量的功能房间"使用MoE'
    })
    
    # 问题3: 专家类型选择不当
    issues.append({
        'issue': '专家类型选择不当',
        'description': 'MoE专家应该是Bottleneck模块，而不是简单的Conv',
        'current': 'Expert使用简单的Conv层',
        'correct': 'Expert应该使用Bottleneck结构',
        'explanation': 'GUIDE/MOE.md指出专家应该是"多个并行的Bottleneck模块"'
    })
    
    # 问题4: 缺少负载均衡损失集成
    issues.append({
        'issue': '缺少负载均衡损失集成',
        'description': '训练流程中没有集成负载均衡损失',
        'current': '只实现了get_load_balancing_loss方法',
        'correct': '需要在train.py中集成负载均衡损失',
        'explanation': 'GUIDE/MOE.md强调这是"最关键也最容易被忽略的一步"'
    })
    
    for i, issue in enumerate(issues, 1):
        print(f"\n❌ 问题 {i}: {issue['issue']}")
        print(f"   描述: {issue['description']}")
        print(f"   当前: {issue['current']}")
        print(f"   正确: {issue['correct']}")
        print(f"   说明: {issue['explanation']}")
    
    return issues


def generate_corrected_yaml():
    """生成修正后的YAML配置"""
    print(f"\n" + "=" * 60)
    print("🛠️  生成修正后的MoE配置")
    print("=" * 60)
    
    corrected_configs = {
        'yolov5s-moe-lite-corrected.yaml': {
            'description': '修正后的轻量级MoE配置',
            'changes': [
                '只在关键C3模块使用MoE',
                '保留标准Conv层',
                '修正参数顺序',
                '使用Bottleneck专家'
            ],
            'backbone': [
                "[-1, 1, Conv, [64, 6, 2, 2]]",  # 0-P1/2
                "[-1, 1, Conv, [128, 3, 2]]",    # 1-P2/4 (保持标准)
                "[-1, 3, C3, [128]]",            # 2 (保持标准)
                "[-1, 1, Conv, [256, 3, 2]]",    # 3-P3/8 (保持标准)
                "[-1, 6, C3MoE, [256, True, 1, 0.5, 4, 2]]",  # 4 (开始使用MoE)
                "[-1, 1, Conv, [512, 3, 2]]",    # 5-P4/16 (保持标准)
                "[-1, 9, C3MoE, [512, True, 1, 0.5, 6, 2]]",  # 6 (深层MoE)
                "[-1, 1, Conv, [1024, 3, 2]]",   # 7-P5/32 (保持标准)
                "[-1, 3, C3MoE, [1024, True, 1, 0.5, 8, 3]]", # 8 (最深层MoE)
                "[-1, 1, SPPF, [1024, 5]]"       # 9
            ]
        },
        
        'yolov5s-moe-corrected.yaml': {
            'description': '修正后的完整MoE配置',
            'changes': [
                '在所有C3模块使用MoE',
                '保留标准Conv和SPPF层',
                '修正参数顺序',
                '渐进式增加专家数量'
            ],
            'backbone': [
                "[-1, 1, Conv, [64, 6, 2, 2]]",  # 0-P1/2
                "[-1, 1, Conv, [128, 3, 2]]",    # 1-P2/4
                "[-1, 3, C3MoE, [128, True, 1, 0.5, 4, 2]]",  # 2 (MoE-C3)
                "[-1, 1, Conv, [256, 3, 2]]",    # 3-P3/8
                "[-1, 6, C3MoE, [256, True, 1, 0.5, 6, 2]]",  # 4 (更多专家)
                "[-1, 1, Conv, [512, 3, 2]]",    # 5-P4/16
                "[-1, 9, C3MoE, [512, True, 1, 0.5, 8, 3]]",  # 6 (深层专家)
                "[-1, 1, Conv, [1024, 3, 2]]",   # 7-P5/32
                "[-1, 3, C3MoE, [1024, True, 1, 0.5, 8, 4]]", # 8 (最深层)
                "[-1, 1, SPPF, [1024, 5]]"       # 9
            ]
        }
    }
    
    for config_name, config_info in corrected_configs.items():
        print(f"\n📝 {config_name}")
        print(f"   描述: {config_info['description']}")
        print(f"   主要修改:")
        for change in config_info['changes']:
            print(f"   • {change}")
    
    return corrected_configs


def suggest_expert_improvements():
    """建议专家网络改进"""
    print(f"\n" + "=" * 60)
    print("💡 专家网络改进建议")
    print("=" * 60)
    
    suggestions = [
        {
            'title': '使用Bottleneck专家',
            'description': '根据GUIDE/MOE.md，专家应该是Bottleneck模块',
            'code': '''
class BottleneckExpert(nn.Module):
    """Bottleneck专家网络"""
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
        super().__init__()
        c_ = int(c2 * e)
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2
    
    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
'''
        },
        
        {
            'title': '专家多样化',
            'description': '不同专家应该有不同的特性',
            'code': '''
# 专家1: 标准Bottleneck (擅长一般特征)
# 专家2: 大核Bottleneck (擅长大尺度特征)  
# 专家3: 深度可分离Bottleneck (擅长轻量化特征)
# 专家4: 扩张卷积Bottleneck (擅长多尺度特征)
'''
        },
        
        {
            'title': '门控网络优化',
            'description': '门控网络应该更智能',
            'code': '''
class ImprovedGating(nn.Module):
    def __init__(self, c1, num_experts, top_k=2):
        super().__init__()
        self.gate = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(c1, c1 // 4),  # 降维
            nn.ReLU(),
            nn.Linear(c1 // 4, num_experts),
            nn.Softmax(dim=-1)
        )
'''
        }
    ]
    
    for i, suggestion in enumerate(suggestions, 1):
        print(f"\n💡 建议 {i}: {suggestion['title']}")
        print(f"   {suggestion['description']}")
        print(f"   代码示例:")
        print(suggestion['code'])


def main():
    """主函数"""
    print("🔧 MoE架构修复分析")
    print("=" * 80)
    
    # 分析问题
    issues = analyze_moe_issues()
    
    # 生成修正配置
    corrected_configs = generate_corrected_yaml()
    
    # 改进建议
    suggest_expert_improvements()
    
    # 总结
    print(f"\n" + "=" * 80)
    print("📋 修复总结")
    print("=" * 80)
    
    print(f"\n🎯 核心问题:")
    print("1. YAML参数配置错误，导致参数解析混乱")
    print("2. 过度使用MoE，违背了'在关键位置使用'的原则")
    print("3. 专家网络设计不符合文档要求")
    print("4. 缺少训练流程中的负载均衡损失集成")
    
    print(f"\n✅ 修复方案:")
    print("1. 修正YAML配置参数顺序")
    print("2. 只在C3模块使用MoE，保留标准Conv层")
    print("3. 使用Bottleneck作为专家网络")
    print("4. 在train.py中集成负载均衡损失")
    
    print(f"\n🚀 下一步行动:")
    print("1. 修复models/common.py中的Expert类")
    print("2. 重新生成正确的YAML配置文件")
    print("3. 在训练流程中添加负载均衡损失")
    print("4. 测试修复后的MoE模型")
    
    print(f"\n" + "=" * 80)


if __name__ == "__main__":
    main()