097、YOLOv8改进实战EMA指数移动平均与SWA随机权重平均在模型训练中的应用从一次模型抖动说起上个月调一个YOLOv8的工业检测项目模型在验证集上loss曲线震荡得像心电图明明训练集收敛得挺好一到验证就抽风。当时我盯着终端里跳动的mAP数字心里那个烦——0.5到0.7之间反复横跳根本没法确定哪个checkpoint能用。后来翻出之前做分类任务时用过的EMA指数移动平均和SWA随机权重平均试了试效果出奇的好。今天就把这两个在YOLOv8里怎么落地、踩过哪些坑一次性说清楚。EMA给模型权重加个“低通滤波器”EMA的原理说白了就是训练过程中维护一份模型权重的滑动平均值推理时用这份平滑后的权重替代原始权重。为什么有效因为SGD的梯度更新本身带有噪声权重在最优解附近来回震荡EMA相当于给权重序列做了个低通滤波把高频抖动滤掉留下更稳定的“重心”。YOLOv8里怎么插EMAYOLOv8的Ultralytics框架其实自带了一个简易的EMA实现在ultralytics/utils/torch_utils.py里有个ModelEMA类。但说实话那个实现太基础了只支持单卡而且衰减系数写死了。我一般自己重写一个嵌入到训练流程里。核心代码就这几行但细节决定成败classEMA:def__init__(self,model,decay0.9999):self.modelmodel self.decaydecay self.shadow{}self.backup{}# 这里踩过坑必须用deepcopy直接赋值会共享内存forname,paraminmodel.named_parameters():ifparam.requires_grad:self.shadow[name]param.data.clone()defupdate(self,model):# 别这样写直接用self.decay * shadow (1-self.decay) * param# 因为训练初期模型变化大decay应该动态调整decaymin(self.decay,(1self.step)/(10self.step))forname,paraminmodel.named_parameters():ifparam.requires_grad:new_average(1-decay)*param.datadecay*self.shadow[name]self.shadow[name]new_average.clone()self.step1defapply_shadow(self):# 推理前调用把EMA权重赋给模型forname,paraminself.model.named_parameters():ifparam.requires_grad:self.backup[name]param.data.clone()param.dataself.shadow[name]defrestore(self):# 推理后恢复原始权重继续训练forname,paraminself.model.named_parameters():ifparam.requires_grad:param.dataself.backup[name]注意那个动态decay——训练初期模型权重变化剧烈固定decay会让EMA更新太慢跟不上模型变化。我参考了PyTorch官方实现的做法让decay从0.999逐渐增加到0.9999这样前期快速适应后期稳定平滑。集成到YOLOv8训练循环在train.py的_do_train方法里找到权重更新那一段大概在optimizer.step()之后# 原始代码optimizer.step()lr_scheduler.step()# 插入EMA更新ifself.emaisnotNone:self.ema.update(model)然后在每个epoch结束时用EMA权重做验证ifself.emaisnotNone:self.ema.apply_shadow()# 这里做验证计算mAPval_resultsself.validator(model)self.ema.restore()else:val_resultsself.validator(model)有个容易忽略的点验证完一定要restore()否则下一轮训练用的就是EMA权重了梯度更新会乱套。我刚开始就忘了这茬训练了50个epoch发现loss不降反升排查了半天。SWA多个好模型的“民主投票”SWA的思路更激进——它不搞在线平滑而是把训练过程中多个checkpoint的权重直接平均。理论上SGD找到的最优点往往在损失平面的平坦区域边缘多个点的平均能让你更接近平坦区域的中心泛化能力更强。YOLOv8里实现SWASWA需要记录多个checkpoint我一般用循环队列classSWA:def__init__(self,model,swa_start0.75,swa_freq5):self.modelmodel self.swa_startswa_start# 训练进度75%时开始收集self.swa_freqswa_freq# 每5个epoch收集一次self.swa_model{}self.n_models0self.total_epochsNonedefupdate(self,model,epoch,total_epochs):self.total_epochstotal_epochsifepochself.swa_start*total_epochs:return# 还没到收集阶段if(epoch-self.swa_start*total_epochs)%self.swa_freq!0:return# 不是收集点# 这里踩过坑直接累加会导致数值溢出应该用在线平均self.n_models1forname,paraminmodel.named_parameters():ifparam.requires_grad:ifnamenotinself.swa_model:self.swa_model[name]param.data.clone()else:# 在线更新平均new_avg old_avg (param - old_avg) / nself.swa_model[name].add_((param.data-self.swa_model[name])/self.n_models)defapply_swa(self):forname,paraminself.model.named_parameters():ifparam.requires_gradandnameinself.swa_model:param.dataself.swa_model[name].clone()SWA的启动时机很关键。我试过从训练一开始就收集结果平均出来的模型还不如单个checkpoint。经验是等模型基本收敛了再开始一般设在训练进度的70%-80%之后。频率上每5-10个epoch收集一次就够了太频繁了平均出来的权重差异小效果不明显。SWA的BN层问题SWA有个大坑平均后的权重BN层的running_mean和running_statistics是错的。因为BN的统计量依赖于模型在训练时的行为平均后的权重对应的统计量需要重新计算。YOLOv8的BN层不多但影响不小。解决方案是在SWA平均后用一小部分训练数据跑一遍前向传播更新BN统计量defupdate_bn(self,dataloader,device):self.model.train()# 必须切到train模式BN才会更新withtorch.no_grad():fori,(images,targets)inenumerate(dataloader):ifi200:# 200个batch足够breakimagesimages.to(device)_self.model(images)self.model.eval()别用太多数据200个batch左右就够了。我试过用整个训练集耗时太长效果提升微乎其微。EMA vs SWA什么时候用哪个这两个东西不是互斥的可以组合使用。我一般这样搭配训练周期短100 epoch只用EMA。SWA需要足够的收集点短周期里收集不到几个checkpoint平均效果有限。训练周期长200 epochEMA SWA一起上。训练过程中用EMA做验证选checkpoint训练结束后用SWA做最终模型。数据量小1000张优先SWA。小数据集容易过拟合SWA的平坦区域偏好能显著提升泛化。数据量大10万张EMA就够了。大数据集下模型本身已经比较稳定SWA的收益不明显还多花时间。实际效果一个工业检测案例说回开头那个项目——检测电路板上的微小焊点缺陷。原始YOLOv8训练100个epoch验证集mAP0.5在0.65到0.72之间震荡最佳checkpoint是0.72。加上EMA后验证曲线平滑了很多最终mAP稳定在0.74。再叠加上SWA从第75个epoch开始每5个epoch收集一次最终模型mAP达到0.77。别小看这3个点的提升在工业场景里0.77和0.72意味着漏检率降低将近20%。推理速度方面EMA和SWA都不增加额外计算量——因为最终模型只有一个权重是提前算好的。唯一代价是训练时多占一点显存EMA需要存一份shadow weightsSWA需要存一份平均权重但相比模型本身这点开销可以忽略。个人经验别迷信默认参数EMA的decay、SWA的启动时机和频率都得根据你的数据集和训练时长调。我见过有人直接套用ImageNet上的参数结果效果还不如不用。EMA和SWA不能替代学习率衰减它们解决的是权重震荡问题不是收敛问题。该用Cosine Annealing还是得用该做Warmup还是得做。多卡训练时注意同步如果用了DistributedDataParallelEMA的shadow weights需要在所有卡上保持一致。我一般只在rank0的卡上维护EMA然后广播给其他卡。验证时用EMA/SWA保存时也保存原始权重有时候EMA/SWA的权重在验证集上表现好但在实际部署场景里可能不如原始权重。我习惯同时保存两份部署前做A/B测试。SWA的BN更新别偷懒我见过有人直接跳过这一步结果模型在推理时BN统计量不对输出结果全是NaN。这个坑我踩过印象深刻。这两个技巧不是什么高大上的创新但确实是工程落地时性价比极高的trick。尤其是当你把YOLOv8调到瓶颈期发现验证集loss怎么都降不下去的时候试试EMA和SWA往往能给你惊喜。