当前位置：首页 > news >正文

YOLO-Master 的MoE方案分解

news 2026/8/18 13:28:49

之前，进行论文精度。今天看下具体代码

文章目录

1. OptimizedMOEImproved
- 加载模块
- 过程
2. 路由模块 EfficientSpatialRouter
3. 专家 SimpleExpert
实例条件自适应
MoE 剪枝 (MoEPruner)
聚类加权 NMS (CW-NMS)

1. OptimizedMOEImproved

同构专家：通常使用相同的 SimpleExpert，便于并行优化。
引入 Shared Expert (共享专家)：增加了一个始终激活的并行分支。这是现代 MoE 的标配，保证了模型的保底性能，显著提升训练稳定性。

高效空间路由：使用 EfficientSpatialRouter，引入预池化（Pre-pooling）减少路由计算量。

稳定性增强：引入 Z-Loss，防止 Router 输出的 Logits 数值爆炸，进一步稳定训练。
标准化的辅助损失：整合了负载均衡损失（Load Balancing Loss）和 Z-Loss。
初始化策略：对 Router 进行了专门的初始化（高斯分布 std=0.01），防止初期“赢家通吃”。

加载模块

router_type 是 EfficientSpatialRouter

self.routing=EfficientSpatialRouter(in_channels,num_experts,top_k=top_k,noise_std=noise_std)

expert_type 是 SimpleExpert

self.experts.append(SimpleExpert(in_channels,out_channels,**kwargs))

共享专家

self.shared_expert=nn.Sequential(nn.Conv2d(in_channels,out_channels,1,bias=False),nn.BatchNorm2d(out_channels),nn.SiLU(inplace=True))

过程

获得路由信息
1） routing_weights, routing_indices, loss_dict = self.routing(x)

共享专家 shared_out = self.shared_expert(x)

3）topk选择

indices_flat=routing_indices.view(B,adaptive_top_k)weights_flat=routing_weights.view(B,adaptive_top_k)

4）专家计算

# Select input and computeinp=x_input[batch_idx]out=self.experts[i](inp)w=weights_flat[batch_idx,k_idx].view(-1,1,1,1)expert_output.index_add_(0,batch_idx,out*w)

输出

final_output = shared_out + expert_output

2. 路由模块 EfficientSpatialRouter

先降采样再路由。通过 AvgPool 减小特征图尺寸，大幅降低 FLOPs。

self.router=nn.Sequential(nn.Conv2d(in_channels,reduced_channels,3,padding=1,bias=False),nn.BatchNorm2d(reduced_channels),nn.SiLU(inplace=True),nn.Conv2d(reduced_channels,num_experts,1,bias=False),nn.BatchNorm2d(num_experts)# numerical stability)

处理过程

global_logits=torch.mean(out,dim=[2,3])# [B, E]self._process_logits(global_logits,self.noise_std,self.training)

def_process_logits(self,logits:torch.Tensor,noise_std:float,training:bool)->Tuple[torch.Tensor,torch.Tensor,Dict]:"""Unified logic to process logits into Top-K selection."""B=logits.shape[0]# 1) Add noise during training (simplified Gumbel-Softmax trick)iftrainingandnoise_std>0:logits=logits+torch.randn_like(logits)*noise_std# 2) Compute probabilitiesprobs=F.softmax(logits.float(),dim=1).type_as(logits)# 3) Select Top-Ktopk_vals,topk_indices=torch.topk(probs,self.top_k,dim=1)# 4) Normalize weightssum_vals=topk_vals.sum(dim=1,keepdim=True)+1e-6topk_vals=topk_vals/sum_vals# 5) Collect loss-related info (train only)loss_dict={}iftraining:loss_dict['router_logits']=logits loss_dict['router_probs']=probs loss_dict['topk_indices']=topk_indicesreturntopk_vals,topk_indices,loss_dict

3. 专家 SimpleExpert

Conv-BN-SiLU-Conv-BN
标准结构，易于优化。
参数量标准。

classSimpleExpert(nn.Module):def__init__(self,in_channels,out_channels,expand_ratio=2):super().__init__()hidden_dim=int(in_channels*expand_ratio)self.conv=nn.Sequential(nn.Conv2d(in_channels,hidden_dim,1,bias=False),nn.BatchNorm2d(hidden_dim),nn.SiLU(inplace=True),nn.Conv2d(hidden_dim,out_channels,1,bias=False),nn.BatchNorm2d(out_channels))defforward(self,x):returnself.conv(x)defcompute_flops(self,input_shape):returnFlopsUtils.count_conv2d(self.conv,input_shape)