第 11 章 · YOLO-Master——MoE 遇见实时检测

# MoE 的前向传播
def moe_forward(x, experts, router, top_k=2):
    # x: 输入特征
    # experts: N 个专家网络
    # router: 路由网络

    # Step 1: 路由器计算每个专家的权重
    gate_logits = router(x)              # [batch, num_experts]
    gate_probs = softmax(gate_logits)     # 归一化

    # Step 2: 选择 Top-K 个专家
    topk_probs, topk_indices = gate_probs.topk(top_k)
    topk_probs = topk_probs / topk_probs.sum()  # 重新归一化

    # Step 3: 只计算被选中的专家
    output = 0
    for i, (prob, idx) in enumerate(zip(topk_probs, topk_indices)):
        output += prob * experts[idx](x)  # 加权求和

    return output

class MoE_Bottleneck_v1(nn.Module):
    """v1: 直接把 Bottleneck 中的 FFN 换成 MoE"""
    def __init__(self, c, num_experts=4, top_k=2):
        super().__init__()
        self.cv1 = Conv(c, c, 3, 1)           # 第一个卷积不变
        self.experts = nn.ModuleList([
            Conv(c, c, 3, 1) for _ in range(num_experts)
        ])
        self.router = nn.Linear(c, num_experts)  # 路由器
        self.top_k = top_k

    def forward(self, x):
        # 路由决策
        # 对特征图做全局平均池化，得到通道描述符
        gate_input = x.mean(dim=[2, 3])  # [B, C]
        gate_logits = self.router(gate_input)  # [B, num_experts]
        topk_probs, topk_idx = gate_logits.softmax(-1).topk(self.top_k)

        # 只计算被选中的专家
        out = sum(
            topk_probs[:, i].view(-1, 1, 1, 1) * self.experts[topk_idx[:, i]](x)
            for i in range(self.top_k)
        )
        return x + out  # 残差连接

# 负载均衡损失
def load_balance_loss(gate_probs, num_experts):
    """鼓励每个专家被选中的频率接近均匀分布"""
    # gate_probs: [batch, num_experts]

    # 每个专家被选中的频率
    expert_freq = gate_probs.mean(dim=0)  # [num_experts]

    # 目标：每个专家的频率 = 1/num_experts
    target = torch.ones_like(expert_freq) / num_experts

    # 辅助损失 = 实际频率和均匀频率的 KL 散度
    aux_loss = (expert_freq * expert_freq.log() - expert_freq * target.log()).sum()
    return aux_loss

class MoE_Bottleneck_v3(nn.Module):
    """v3: 共享专家 + 路由专家"""
    def __init__(self, c, num_experts=4, num_shared=1, top_k=2):
        super().__init__()
        # 共享专家：始终激活，处理通用特征
        self.shared_experts = nn.ModuleList([
            Conv(c, c, 3, 1) for _ in range(num_shared)
        ])
        # 路由专家：按需激活，处理特定类型的目标
        self.routed_experts = nn.ModuleList([
            Conv(c, c, 3, 1) for _ in range(num_experts)
        ])
        self.router = nn.Linear(c, num_experts)
        self.top_k = top_k

    def forward(self, x):
        # 共享专家的输出（始终计算）
        shared_out = sum(exp(x) for exp in self.shared_experts)

        # 路由专家的输出（按需计算）
        gate_input = x.mean(dim=[2, 3])
        gate_logits = self.router(gate_input)
        topk_probs, topk_idx = gate_logits.softmax(-1).topk(self.top_k)

        routed_out = sum(
            topk_probs[:, i].view(-1, 1, 1, 1) * self.routed_experts[topk_idx[:, i]](x)
            for i in range(self.top_k)
        )

        return x + shared_out + routed_out

# Z-Loss: 防止路由器 logits 过大导致不稳定
def z_loss(gate_logits):
    """惩罚路由器输出的绝对值过大"""
    # gate_logits: [batch, num_experts]
    # 如果 logits 的绝对值很大，softmax 会趋近于 one-hot
    # 导致梯度消失 + 路由震荡
    return (gate_logits.logsumexp(dim=-1) ** 2).mean()

class ChannelRouter(nn.Module):
    """基于通道统计的路由器"""
    def __init__(self, in_channels, num_experts):
        super().__init__()
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),    # 全局平均池化
            nn.Flatten(),
            nn.Linear(in_channels, num_experts),  # 输出专家权重
        )

    def forward(self, x):
        return self.fc(x)