097、YOLOv8改进实战:EMA指数移动平均与SWA随机权重平均在模型训练中的应用
097、YOLOv8改进实战:EMA指数移动平均与SWA随机权重平均在模型训练中的应用
从一次模型抖动说起
上个月调一个YOLOv8的工业检测项目,模型在验证集上loss曲线震荡得像心电图,明明训练集收敛得挺好,一到验证就抽风。当时我盯着终端里跳动的mAP数字,心里那个烦——0.5到0.7之间反复横跳,根本没法确定哪个checkpoint能用。后来翻出之前做分类任务时用过的EMA(指数移动平均)和SWA(随机权重平均),试了试,效果出奇的好。今天就把这两个在YOLOv8里怎么落地、踩过哪些坑,一次性说清楚。
EMA:给模型权重加个“低通滤波器”
EMA的原理说白了就是:训练过程中,维护一份模型权重的滑动平均值,推理时用这份平滑后的权重替代原始权重。为什么有效?因为SGD的梯度更新本身带有噪声,权重在最优解附近来回震荡,EMA相当于给权重序列做了个低通滤波,把高频抖动滤掉,留下更稳定的“重心”。
YOLOv8里怎么插EMA
YOLOv8的Ultralytics框架其实自带了一个简易的EMA实现,在ultralytics/utils/torch_utils.py里有个ModelEMA类。但说实话,那个实现太基础了,只支持单卡,而且衰减系数写死了。我一般自己重写一个,嵌入到训练流程里。
核心代码就这几行,但细节决定成败:
classEMA:def__init__(self,model,decay=0.9999):self.model=model self.decay=decay 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应该动态调整decay=min(self.decay,(1+self.step)/(10+self.step))forname,paraminmodel.named_parameters():ifparam.requires_grad:new_average=(1-decay)*param.data+decay*self.shadow[name]self.shadow[name]=new_average.clone()self.step+=1defapply_shadow(self):# 推理前调用,把EMA权重赋给模型forname,paraminself.model.named_parameters():ifparam.requires_grad:self.backup[name]=param.data.clone()param.data=self.shadow[name]defrestore(self):# 推理后恢复原始权重,继续训练forname,paraminself.model.named_parameters():ifparam.requires_grad:param.data=self.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_results=self.validator(model)self.ema.restore()else:val_results=self.validator(model)有个容易忽略的点:验证完一定要restore(),否则下一轮训练用的就是EMA权重了,梯度更新会乱套。我刚开始就忘了这茬,训练了50个epoch发现loss不降反升,排查了半天。
SWA:多个好模型的“民主投票”
SWA的思路更激进——它不搞在线平滑,而是把训练过程中多个checkpoint的权重直接平均。理论上,SGD找到的最优点往往在损失平面的平坦区域边缘,多个点的平均能让你更接近平坦区域的中心,泛化能力更强。
YOLOv8里实现SWA
SWA需要记录多个checkpoint,我一般用循环队列:
classSWA:def__init__(self,model,swa_start=0.75,swa_freq=5):self.model=model self.swa_start=swa_start# 训练进度75%时开始收集self.swa_freq=swa_freq# 每5个epoch收集一次self.swa_model={}self.n_models=0self.total_epochs=Nonedefupdate(self,model,epoch,total_epochs):self.total_epochs=total_epochsifepoch<self.swa_start*total_epochs:return# 还没到收集阶段if(epoch-self.swa_start*total_epochs)%self.swa_freq!=0:return# 不是收集点# 这里踩过坑:直接累加会导致数值溢出,应该用在线平均self.n_models+=1forname,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.data=self.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):ifi>200:# 200个batch足够breakimages=images.to(device)_=self.model(images)self.model.eval()别用太多数据,200个batch左右就够了。我试过用整个训练集,耗时太长,效果提升微乎其微。
EMA vs SWA:什么时候用哪个
这两个东西不是互斥的,可以组合使用。我一般这样搭配:
- 训练周期短(<100 epoch):只用EMA。SWA需要足够的收集点,短周期里收集不到几个checkpoint,平均效果有限。
- 训练周期长(>200 epoch):EMA + SWA一起上。训练过程中用EMA做验证选checkpoint,训练结束后用SWA做最终模型。
- 数据量小(<1000张):优先SWA。小数据集容易过拟合,SWA的平坦区域偏好能显著提升泛化。
- 数据量大(>10万张):EMA就够了。大数据集下模型本身已经比较稳定,SWA的收益不明显,还多花时间。
实际效果:一个工业检测案例
说回开头那个项目——检测电路板上的微小焊点缺陷。原始YOLOv8训练100个epoch,验证集mAP@0.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 weights,SWA需要存一份平均权重),但相比模型本身,这点开销可以忽略。
个人经验
别迷信默认参数:EMA的decay、SWA的启动时机和频率,都得根据你的数据集和训练时长调。我见过有人直接套用ImageNet上的参数,结果效果还不如不用。
EMA和SWA不能替代学习率衰减:它们解决的是权重震荡问题,不是收敛问题。该用Cosine Annealing还是得用,该做Warmup还是得做。
多卡训练时注意同步:如果用了DistributedDataParallel,EMA的shadow weights需要在所有卡上保持一致。我一般只在rank=0的卡上维护EMA,然后广播给其他卡。
验证时用EMA/SWA,保存时也保存原始权重:有时候EMA/SWA的权重在验证集上表现好,但在实际部署场景里可能不如原始权重。我习惯同时保存两份,部署前做A/B测试。
SWA的BN更新别偷懒:我见过有人直接跳过这一步,结果模型在推理时BN统计量不对,输出结果全是NaN。这个坑我踩过,印象深刻。
这两个技巧不是什么高大上的创新,但确实是工程落地时性价比极高的trick。尤其是当你把YOLOv8调到瓶颈期,发现验证集loss怎么都降不下去的时候,试试EMA和SWA,往往能给你惊喜。
