1. 从一次模型训练中的“诡异”现象说起最近在复现一个图像分类项目时我遇到了一个挺有意思的现象在训练集上我的模型准确率一路高歌猛进很快就冲到了99%以上损失也降得非常低。但当我满怀信心地把模型拿到验证集上一测结果却让人大跌眼镜——准确率只有70%出头损失也高得离谱。这典型的过拟合现象让我开始排查原因。在检查了数据增强、模型复杂度、正则化等一系列常规操作后一个被我长期忽略的细节浮出水面模型参数的更新方式。我使用的是最基础的随机梯度下降SGD优化器每次迭代都直接用计算出的梯度更新权重。这种“即时生效”的更新方式会让模型权重在训练过程中剧烈波动尤其是在训练后期当学习率还比较大或者遇到一些噪声较大的批次时权重的“瞬时值”可能并不代表其“真实水平”。这就好比一个运动员某一次测试成绩特别好但这并不能代表他的稳定实力。我需要一个能反映模型权重“长期稳定水平”的指标而不是被最后一次更新“带偏”的瞬时值。这时一个在深度学习中看似不起眼实则至关重要的技术进入了我的视野——指数移动平均。指数移动平均英文全称Exponential Moving Average简称EMA。它不是一个独立的优化算法而是一种在模型训练过程中对模型参数进行“平滑”处理的技巧。其核心思想是不直接使用模型在每次迭代后更新得到的瞬时权重而是维护一个权重的“影子”版本。这个影子权重是历史所有瞬时权重的加权平均并且越近的权重占比越高。最终在模型评估或推理时我们使用这个更平滑、更稳定的影子权重而不是最后一次迭代的“毛刺”权重。这个简单的操作往往能带来模型泛化能力的显著提升尤其是在计算机视觉、自然语言处理等对模型稳定性要求较高的领域。2. EMA的核心原理为什么“平均”比“瞬时”更可靠要理解EMA为什么有效我们需要先抛开公式从直观感受和数学本质两个层面来剖析。2.1 直观理解滤除噪声捕捉趋势想象一下股票价格的K线图。如果只看每分钟的股价跳动那曲线会非常“毛糙”充满了各种随机的买卖单造成的瞬时波动。这种波动就是“噪声”它掩盖了股票真正的长期趋势。为了看清趋势分析师们会引入各种移动平均线比如5日均线、20日均线。EMA就是一种特殊的移动平均它给近期的数据点赋予更高的权重因此对价格变化的反应比简单移动平均更灵敏同时又平滑掉了大部分瞬时噪声。在模型训练中情况高度相似。每一次基于一个小批次mini-batch计算出的梯度都像是股价的“分钟线”。这个小批次的数据分布可能并不能完美代表整个数据集因此基于它计算出的梯度更新方向也带有一定的“噪声”。直接用这个带噪声的梯度去更新权重就会让权重在最优值附近来回震荡而不是稳定地趋近于它。EMA所做的就是为模型的每一个权重参数都绘制一条“移动平均线”。这条线滤除了单次更新带来的剧烈波动保留了权重变化的整体趋势使得最终用于推理的权重是一个更接近“真实平均水平”的稳定值。2.2 数学本质递推公式与衰减系数EMA的数学定义非常简洁优雅。假设在训练的第t步我们模型当前的权重为θ_t我们维护的影子权重即EMA权重为θ_t。EMA的更新规则如下θ_t decay * θ_{t-1} (1 - decay) * θ_t这里decay是一个介于0和1之间的超参数通常非常接近1例如0.999, 0.9999。我们更常用另一个参数α 1 - decay它被称为平滑因子或动量系数。那么公式可以重写为θ_t (1 - α) * θ_{t-1} α * θ_t这个递推公式是理解EMA的关键。让我们展开来看当t1时θ_1 α * θ_1假设初始影子权重θ_0 0当t2时θ_2 (1-α)*θ_1 α*θ_2 α*θ_2 α*(1-α)*θ_1当t3时θ_3 α*θ_3 α*(1-α)*θ_2 α*(1-α)^2*θ_1...以此类推你会发现第t步的影子权重θ_t实际上是历史上所有权重θ_1, θ_2, ..., θ_t的加权和其中第k步权重θ_k的系数是α * (1-α)^(t-k)。由于(1-α)小于1所以这个系数随着k的变小即时间越早而呈指数衰减。这就是“指数移动平均”名称的由来——历史权重的影响以指数形式衰减。参数α的选择是门艺术α较大如0.1对应decay0.9影子权重对近期变化非常敏感平滑效果弱更接近瞬时权重。α较小如0.001对应decay0.999影子权重变化非常缓慢平滑效果强能有效滤除噪声但对权重最新趋势的反应也变迟钝。在深度学习实践中decay通常设置为0.999、0.9995或0.9999这意味着α非常小0.001到0.0001。这样的设置保证了影子权重是一个对成百上千次迭代结果进行平滑的稳定值。2.3 与普通SGD和动量的区别这里容易产生混淆我特别说明一下EMA与SGD、SGD with Momentum的区别普通SGDθ_t θ_{t-1} - η * g_t。直接使用当前梯度g_t更新权重。SGD with Momentumv_t β * v_{t-1} g_t;θ_t θ_{t-1} - η * v_t。它是对梯度进行移动平均v_t是平均梯度然后用平均梯度去更新权重。权重θ_t本身仍然是“瞬时值”。EMAθ_t decay * θ_{t-1} (1 - decay) * θ_t。它是对权重本身进行移动平均产生一个独立的影子权重θ_t。训练时模型权重θ_t仍按原方式可以是SGD、Adam等更新但EMA平行地维护另一套平滑后的权重。关键区别在于操作对象动量平均的是梯度目的是让优化方向更稳定EMA平均的是权重目的是得到一个用于推理的更平滑、泛化更好的模型参数。两者可以同时使用并不冲突。3. EMA在PyTorch中的两种实现与避坑指南理论懂了关键还得能落地。在PyTorch中实现EMA我实践下来主要有两种主流方式各有优劣和坑点。3.1 方式一手动维护影子张量灵活但繁琐这是最直接的方式你需要为模型中每一个需要平滑的参数注册一个对应的影子张量。import torch import torch.nn as nn class ModelEMA: def __init__(self, model, decay0.999): self.model model self.decay decay self.shadow {} self.backup {} # 用于临时保存原始权重 # 初始化影子权重 for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] param.data.clone() def update(self): 在每次模型权重更新后调用此方法更新影子权重 for name, param in self.model.named_parameters(): if param.requires_grad: assert name in self.shadow new_average (1.0 - self.decay) * param.data self.decay * self.shadow[name] self.shadow[name] new_average.clone() # 必须使用.clone()避免引用 def apply_shadow(self): 在验证/测试前调用将影子权重应用到模型 for name, param in self.model.named_parameters(): if param.requires_grad: self.backup[name] param.data.clone() param.data self.shadow[name] def restore(self): 在验证/测试后调用恢复模型的原始权重以便继续训练 for name, param in self.model.named_parameters(): if param.requires_grad: param.data self.backup[name] self.backup.clear()使用方式model MyModel() ema ModelEMA(model, decay0.999) # 训练循环 for epoch in range(num_epochs): for data, target in train_loader: # ... 前向传播计算损失 loss.backward() optimizer.step() optimizer.zero_grad() # 关键步骤更新EMA影子权重 ema.update() # 验证阶段 ema.apply_shadow() # 应用EMA权重 evaluate(model, val_loader) # 使用平滑后的权重进行评估 ema.restore() # 恢复原始权重继续训练避坑点1clone()的必要性注意update和apply_shadow方法中我们对张量都使用了.clone()。这是至关重要的。在PyTorch中直接赋值self.shadow[name] new_average会导致self.shadow[name]和new_average共享同一块内存。随后new_average被释放或修改会意外地改变影子权重。.clone()创建了一个数据的独立副本确保了影子权重的独立性。避坑点2仅对可训练参数操作我们通过param.requires_grad进行判断。对于固定不变的参数如预训练模型中被冻结的层或者像BatchNorm的running_mean/running_var这种在训练中通过移动平均更新的统计量通常不应对其再做EMA。对统计量做EMA会导致其更新规则被干扰可能影响模型性能。避坑点3初始化的影响在训练初期模型权重变化剧烈且影子权重是从初始权重开始平均的。如果一开始就应用EMA可能会将模型“拉回”到较差的初始点。一个常见的技巧是设置一个预热步数warmup steps在预热期内让decay从一个较小的值如0.9线性或余弦增长到目标值如0.999让EMA在训练初期更快地跟上权重的变化。3.2 方式二使用PyTorch内置的torch.optim.swa_utils官方推荐但需注意细节从PyTorch 1.6开始官方在torch.optim.swa_utils中提供了对随机权重平均SWA的支持而EMA可以看作是SWA的一种特例使用固定衰减系数。AveragedModel类本质上就是一个EMA。import torch.optim.swa_utils as swa_utils model MyModel() # 创建EMA模型avg_fn参数指定了EMA的更新规则 ema_model swa_utils.AveragedModel(model, avg_fnlambda averaged_model_parameter, model_parameter, num_averaged: decay * averaged_model_parameter (1 - decay) * model_parameter) # 训练循环中在optimizer.step()之后 ema_model.update_parameters(model) # 验证时直接使用ema_model with torch.no_grad(): for data, target in val_loader: output ema_model(data) # ... 计算指标优点代码简洁官方维护自动处理设备CPU/GPU和张量拷贝问题。坑点默认情况下AveragedModel会平均模型的所有参数。如果你的模型包含BatchNorm层这可能会出问题。因为BatchNorm层除了权重weight和偏置bias还有在训练中动态更新的running_mean和running_var。对这些统计量做平均是不合适的。解决方案使用swa_utils.get_ema_avg_fn或自定义avg_fn并在创建AveragedModel时传入一个device参数同时更稳妥的做法是在验证前调用swa_utils.update_bn函数基于训练数据重新计算EMA模型的BatchNorm统计量。但这需要额外遍历一遍训练数据增加开销。# 更安全的做法在训练结束后用训练数据更新EMA模型的BN统计量 ema_model swa_utils.AveragedModel(model) # ... 训练循环不断调用 ema_model.update_parameters(model) # 训练结束后 swa_utils.update_bn(train_loader, ema_model, devicedevice) # 更新BN统计量 # 然后再进行最终验证或保存模型个人经验对于研究或快速实验手动实现的ModelEMA类给了我更大的灵活性和可控性特别是对于复杂模型或需要特殊处理如部分参数冻结的情况。而对于生产环境或标准流程使用官方的AveragedModel并妥善处理BN层是更干净、更不易出错的选择。4. EMA超参数调优与效果验证实战设置好EMA代码只是第一步让它真正发挥作用还需要仔细调整超参数并设计实验来验证效果。4.1 核心超参数decay的调优策略decay是EMA唯一的超参数但它对最终模型性能的影响非常显著。我的调优经验遵循以下步骤确定范围对于大多数视觉和NLP任务decay的有效范围通常在[0.99, 0.9999]之间。可以从一个中间值开始比如0.999。结合训练周期decay的选择与总训练迭代次数T密切相关。一个经验法则是希望EMA权重能覆盖足够多的历史迭代。影子权重中最近N ≈ 1 / (1 - decay)次迭代的贡献占主导。例如decay0.99-N≈100次迭代decay0.999-N≈1000次迭代decay0.9999-N≈10000次迭代 你的N应该远小于总迭代次数T否则EMA权重会被早期不成熟的权重过度影响。如果总共只训练5000步用decay0.9999(N10000) 就太大了。网格搜索与观察曲线在一个小的验证集上对几个候选值如0.995, 0.999, 0.9995进行网格搜索。不仅要看最终的验证准确率更要观察验证损失曲线。一个合适的decay应该能使验证损失曲线更平滑、下降更稳定并且最终稳定在一个更低的平台。考虑预热Warmup如前所述在训练初期使用较低的decay或较高的α有助于EMA权重快速跟上模型权重的变化。你可以实现一个动态的decaydef get_current_decay(iter, warmup_iters1000, base_decay0.999): if iter warmup_iters: # 线性从0.9增长到base_decay return 0.9 (base_decay - 0.9) * (iter / warmup_iters) else: return base_decay在每次update时计算当前decay并传入。4.2 效果验证不仅仅是看最终准确率验证EMA是否有效不能只看最终验证集上的一个准确率数字。你需要进行更细致的分析训练vs验证损失曲线对比这是最直观的。在同一张图上绘制使用原始权重和EMA权重计算出的验证损失曲线。一个成功的EMA应该能降低验证损失的波动曲线更平滑。降低验证损失的最终值曲线收敛到更低的点。缓解过拟合训练损失和验证损失之间的差距缩小。权重分布可视化在训练的不同阶段早期、中期、后期分别提取原始权重和EMA权重中某一层的参数如全连接层的权重绘制其直方图或计算其统计量均值、方差。你会发现EMA权重的分布通常更加“集中”方差更小极端值过大或过小的权重更少。这从理论上解释了其更好的泛化性复杂的、过拟合的模型往往拥有一些绝对值非常大的权重EMA平滑了这些极端值。鲁棒性测试对验证集数据施加轻微的扰动如高斯噪声、轻微的模糊、色彩抖动然后分别用原始模型和EMA模型进行测试。EMA模型在扰动下的性能下降通常更小表现出更强的鲁棒性。“快照”集成效应由于EMA权重是历史权重的平均它在某种程度上近似于将训练过程中多个时间点的模型“快照”进行了集成。你可以做一个对比实验单独保存训练过程中几个检查点checkpoint在推理时对它们的预测结果进行平均软投票。这个结果与单一EMA模型的结果进行对比。在很多情况下EMA模型能达到甚至超过多模型集成的效果却只需要存储和运行一个模型效率极高。4.3 一个完整的训练脚本片段示例结合以上所有要点一个整合了动态decay、验证和保存的EMA训练循环核心部分如下import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader import copy class EMA: # ... 使用前面手动实现的ModelEMA类但增加动态decay def __init__(self, model, base_decay0.999, warmup_iters1000): self.model model self.base_decay base_decay self.warmup_iters warmup_iters self.shadow {} self.backup {} for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] param.data.clone().detach() def get_decay(self, iter): if iter self.warmup_iters: return 0.9 (self.base_decay - 0.9) * min(1.0, iter / self.warmup_iters) return self.base_decay def update(self, iter): decay self.get_decay(iter) for name, param in self.model.named_parameters(): if param.requires_grad: self.shadow[name].data.copy_(decay * self.shadow[name].data (1.0 - decay) * param.data) # ... apply_shadow, restore 方法同上 # 初始化 model MyModel().cuda() ema EMA(model, base_decay0.999, warmup_iters2000) optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() best_val_acc 0.0 global_iter 0 for epoch in range(100): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target data.cuda(), target.cuda() optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() # 更新EMA ema.update(global_iter) global_iter 1 # 验证阶段 model.eval() # 使用原始模型验证 orig_val_acc evaluate(model, val_loader) # 使用EMA模型验证 ema.apply_shadow() ema_val_acc evaluate(model, val_loader) # 此时model的权重已被替换为EMA权重 ema.restore() print(fEpoch {epoch}: Orig Val Acc: {orig_val_acc:.4f}, EMA Val Acc: {ema_val_acc:.4f}) # 保存最佳EMA模型 if ema_val_acc best_val_acc: best_val_acc ema_val_acc ema.apply_shadow() best_model_state copy.deepcopy(model.state_dict()) # 保存的是EMA权重 ema.restore() torch.save({ epoch: epoch, model_state_dict: best_model_state, ema_shadow: ema.shadow, # 也可以选择保存影子字典 optimizer_state_dict: optimizer.state_dict(), best_acc: best_val_acc, }, best_ema_model.pth)这个流程清晰地展示了EMA如何无缝嵌入训练循环并通过对比验证准确率来体现其价值。保存模型时务必注意你保存的是应用了影子权重后的模型状态best_model_state还是影子字典本身。前者加载后可直接用于推理后者则需要先初始化EMA类再加载。5. 进阶话题EMA的变体、局限性与相关技术掌握了基础用法后我们可以看看EMA的一些高级变体、它不适用的场景以及与其相关的其他技术。5.1 EMA的常见变体带偏置校正的EMA在训练的最初几步由于影子权重从零或初始值开始其值会偏向于初始值存在偏差。尤其是在decay接近1时这个偏差在早期会很明显。偏置校正通过在早期对EMA值进行缩放来消除这个影响。公式为θ_t_corrected θ_t / (1 - decay^t)。这在Adam等优化器中很常见但在模型权重的EMA中由于我们通常关心训练稳定后的最终权重且预热策略也能缓解此问题所以不常使用。周期性EMAStochastic Weight Averaging, SWASWA可以看作是EMA的一个“节拍器”变体。它不是在每一步都更新影子权重而是以固定的周期如每几个epoch将当前权重加入到平均池中。更新公式是简单的算术平均θ_new (θ_old * n θ_current) / (n 1)。SWA的理论基础是SGD优化路径会在最优解周围的多模盆地中游走平均这些权重可以落在更中心的泛化更好的区域。SWA通常不需要调整decay这样的超参数且在许多任务上表现比固定decay的EMA更鲁棒。PyTorch的swa_utils主要就是为SWA设计的。EMA与学习率调度的协同当使用学习率衰减策略时如StepLR、CosineAnnealing模型权重在后期更新幅度变小。此时EMA的平滑效应会更强。一个有趣的实践是在训练末期可以增大decay例如从0.999增加到0.9999让影子权重变化更慢进一步平滑最后阶段的微小波动有助于模型收敛到更平坦的极小值这通常与更好的泛化性相关。5.2 EMA的局限性什么时候可能没用甚至有害EMA不是银弹在以下场景需要谨慎使用或避免使用训练数据极度干净过拟合风险极低如果你的模型很简单或者数据量极大模型本身就不容易过拟合那么EMA带来的提升可能微乎其微白增加了复杂性。优化器本身已具备强平滑性如果你使用的是像AdamW这样自适应学习率且自带动量对梯度一阶矩估计的优化器它已经在一定程度上平滑了更新过程。再加上EMA效果可能叠加也可能过度平滑导致模型收敛变慢。需要实验验证。对BatchNorm层处理不当这是最大的坑。如前所述对BN层的running_mean/var做EMA会破坏其统计特性。标准做法是EMA只应用于模型的可学习参数权重和偏置而不应用于BN的统计量。在验证/测试时应使用EMA权重但BN层应使用其自身在训练中累积的running_mean/var即原始模型BN层的统计量或者最好在训练结束后用EMA权重重新前向传播一遍训练数据来更新BN统计量update_bn。动态网络结构或权重共享对于结构在训练中发生变化的网络如某些NAS方法或权重共享的模块EMA的更新逻辑可能变得复杂需要特别设计。训练初期在模型权重快速变化的初期阶段过早应用强力的EMA高decay会拖慢影子权重的更新可能不利于模型快速找到有希望的区域。这就是为什么预热策略很重要。5.3 与EMA相关的热词EMA注意力机制最近在搜索EMA时常会看到“EMA注意力机制”这个词。这里需要做一个重要的区分本文讨论的EMA指数移动平均是一种应用于模型参数的平滑技术用于提升模型泛化能力。它是一个训练技巧。EMA注意力机制Efficient Multi-scale Attention这是一种网络模块结构的设计通常出现在计算机视觉的骨干网络中如EMANet Efficient Multi-scale Attention Network。它通过引入多尺度上下文信息和高效的注意力计算来提升模型的特征提取能力。这里的“EMA”是模块的名称缩写其内部可能使用了移动平均的思想来进行特征融合但和本文所述的参数平滑技术是完全不同的概念和应用层面。不要混淆两者。当你看到一篇论文或代码中提到“EMA模块”时需要根据上下文判断它指的是参数平滑技巧还是一种特定的神经网络层。6. 总结与个人心得回顾整个探索过程EMA给我的最大启示是在追求模型性能的路上我们不仅要关注“前沿”和“复杂”更要重视那些被验证有效的“基础”与“简单”。EMA几乎没有增加任何计算开销只是多存储一套权重和一些简单的标量运算却能稳定地带来1-2个百分点的泛化性能提升这在很多竞赛和实际项目中可能就是决定性的优势。从我个人的实战经验来看以下几点心得或许对你有帮助把它变成默认选项对于大多数监督学习任务尤其是视觉和NLP任务我现在会习惯性地在训练脚本里加上EMA。它的收益风险比极高。先跑基线再加EMA在调试新模型或新任务时我通常会先不用EMA跑一个基线观察训练和验证曲线的正常形态。然后再加入EMA对比曲线变化这样能更清晰地看到EMA的效果并帮助我设置合适的decay和预热策略。保存与加载的细节决定成败多少次因为保存和加载EMA模型的状态不对而debug到深夜。务必明确你保存的是什么原始权重、影子字典还是应用了影子的模型状态并在加载时进行对应的恢复操作。写一个清晰的文档字符串或注释说明保存格式至关重要。不要神话它EMA是优秀的正则化工具但它不能替代良好的数据、合适的模型结构和正确的优化器选择。它是“锦上添花”而非“雪中送炭”。如果模型在基础设定下都无法收敛先别指望EMA能拯救世界。最后技术总是在发展。EMA和SWA这类权重平均技术本质上是在损失函数的权重空间里寻找更平坦、泛化更好的解。这与当前关于“平坦极小值”与泛化性的理论研究是相呼应的。理解其背后的思想比单纯调用一个API更有价值。下次当你训练模型时不妨花几分钟加上这几行代码它可能会给你带来意想不到的回报。 深度学习模型泛化提升利器：指数移动平均（EMA）原理与PyTorch实战
1. 从一次模型训练中的“诡异”现象说起最近在复现一个图像分类项目时我遇到了一个挺有意思的现象在训练集上我的模型准确率一路高歌猛进很快就冲到了99%以上损失也降得非常低。但当我满怀信心地把模型拿到验证集上一测结果却让人大跌眼镜——准确率只有70%出头损失也高得离谱。这典型的过拟合现象让我开始排查原因。在检查了数据增强、模型复杂度、正则化等一系列常规操作后一个被我长期忽略的细节浮出水面模型参数的更新方式。我使用的是最基础的随机梯度下降SGD优化器每次迭代都直接用计算出的梯度更新权重。这种“即时生效”的更新方式会让模型权重在训练过程中剧烈波动尤其是在训练后期当学习率还比较大或者遇到一些噪声较大的批次时权重的“瞬时值”可能并不代表其“真实水平”。这就好比一个运动员某一次测试成绩特别好但这并不能代表他的稳定实力。我需要一个能反映模型权重“长期稳定水平”的指标而不是被最后一次更新“带偏”的瞬时值。这时一个在深度学习中看似不起眼实则至关重要的技术进入了我的视野——指数移动平均。指数移动平均英文全称Exponential Moving Average简称EMA。它不是一个独立的优化算法而是一种在模型训练过程中对模型参数进行“平滑”处理的技巧。其核心思想是不直接使用模型在每次迭代后更新得到的瞬时权重而是维护一个权重的“影子”版本。这个影子权重是历史所有瞬时权重的加权平均并且越近的权重占比越高。最终在模型评估或推理时我们使用这个更平滑、更稳定的影子权重而不是最后一次迭代的“毛刺”权重。这个简单的操作往往能带来模型泛化能力的显著提升尤其是在计算机视觉、自然语言处理等对模型稳定性要求较高的领域。2. EMA的核心原理为什么“平均”比“瞬时”更可靠要理解EMA为什么有效我们需要先抛开公式从直观感受和数学本质两个层面来剖析。2.1 直观理解滤除噪声捕捉趋势想象一下股票价格的K线图。如果只看每分钟的股价跳动那曲线会非常“毛糙”充满了各种随机的买卖单造成的瞬时波动。这种波动就是“噪声”它掩盖了股票真正的长期趋势。为了看清趋势分析师们会引入各种移动平均线比如5日均线、20日均线。EMA就是一种特殊的移动平均它给近期的数据点赋予更高的权重因此对价格变化的反应比简单移动平均更灵敏同时又平滑掉了大部分瞬时噪声。在模型训练中情况高度相似。每一次基于一个小批次mini-batch计算出的梯度都像是股价的“分钟线”。这个小批次的数据分布可能并不能完美代表整个数据集因此基于它计算出的梯度更新方向也带有一定的“噪声”。直接用这个带噪声的梯度去更新权重就会让权重在最优值附近来回震荡而不是稳定地趋近于它。EMA所做的就是为模型的每一个权重参数都绘制一条“移动平均线”。这条线滤除了单次更新带来的剧烈波动保留了权重变化的整体趋势使得最终用于推理的权重是一个更接近“真实平均水平”的稳定值。2.2 数学本质递推公式与衰减系数EMA的数学定义非常简洁优雅。假设在训练的第t步我们模型当前的权重为θ_t我们维护的影子权重即EMA权重为θ_t。EMA的更新规则如下θ_t decay * θ_{t-1} (1 - decay) * θ_t这里decay是一个介于0和1之间的超参数通常非常接近1例如0.999, 0.9999。我们更常用另一个参数α 1 - decay它被称为平滑因子或动量系数。那么公式可以重写为θ_t (1 - α) * θ_{t-1} α * θ_t这个递推公式是理解EMA的关键。让我们展开来看当t1时θ_1 α * θ_1假设初始影子权重θ_0 0当t2时θ_2 (1-α)*θ_1 α*θ_2 α*θ_2 α*(1-α)*θ_1当t3时θ_3 α*θ_3 α*(1-α)*θ_2 α*(1-α)^2*θ_1...以此类推你会发现第t步的影子权重θ_t实际上是历史上所有权重θ_1, θ_2, ..., θ_t的加权和其中第k步权重θ_k的系数是α * (1-α)^(t-k)。由于(1-α)小于1所以这个系数随着k的变小即时间越早而呈指数衰减。这就是“指数移动平均”名称的由来——历史权重的影响以指数形式衰减。参数α的选择是门艺术α较大如0.1对应decay0.9影子权重对近期变化非常敏感平滑效果弱更接近瞬时权重。α较小如0.001对应decay0.999影子权重变化非常缓慢平滑效果强能有效滤除噪声但对权重最新趋势的反应也变迟钝。在深度学习实践中decay通常设置为0.999、0.9995或0.9999这意味着α非常小0.001到0.0001。这样的设置保证了影子权重是一个对成百上千次迭代结果进行平滑的稳定值。2.3 与普通SGD和动量的区别这里容易产生混淆我特别说明一下EMA与SGD、SGD with Momentum的区别普通SGDθ_t θ_{t-1} - η * g_t。直接使用当前梯度g_t更新权重。SGD with Momentumv_t β * v_{t-1} g_t;θ_t θ_{t-1} - η * v_t。它是对梯度进行移动平均v_t是平均梯度然后用平均梯度去更新权重。权重θ_t本身仍然是“瞬时值”。EMAθ_t decay * θ_{t-1} (1 - decay) * θ_t。它是对权重本身进行移动平均产生一个独立的影子权重θ_t。训练时模型权重θ_t仍按原方式可以是SGD、Adam等更新但EMA平行地维护另一套平滑后的权重。关键区别在于操作对象动量平均的是梯度目的是让优化方向更稳定EMA平均的是权重目的是得到一个用于推理的更平滑、泛化更好的模型参数。两者可以同时使用并不冲突。3. EMA在PyTorch中的两种实现与避坑指南理论懂了关键还得能落地。在PyTorch中实现EMA我实践下来主要有两种主流方式各有优劣和坑点。3.1 方式一手动维护影子张量灵活但繁琐这是最直接的方式你需要为模型中每一个需要平滑的参数注册一个对应的影子张量。import torch import torch.nn as nn class ModelEMA: def __init__(self, model, decay0.999): self.model model self.decay decay self.shadow {} self.backup {} # 用于临时保存原始权重 # 初始化影子权重 for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] param.data.clone() def update(self): 在每次模型权重更新后调用此方法更新影子权重 for name, param in self.model.named_parameters(): if param.requires_grad: assert name in self.shadow new_average (1.0 - self.decay) * param.data self.decay * self.shadow[name] self.shadow[name] new_average.clone() # 必须使用.clone()避免引用 def apply_shadow(self): 在验证/测试前调用将影子权重应用到模型 for name, param in self.model.named_parameters(): if param.requires_grad: self.backup[name] param.data.clone() param.data self.shadow[name] def restore(self): 在验证/测试后调用恢复模型的原始权重以便继续训练 for name, param in self.model.named_parameters(): if param.requires_grad: param.data self.backup[name] self.backup.clear()使用方式model MyModel() ema ModelEMA(model, decay0.999) # 训练循环 for epoch in range(num_epochs): for data, target in train_loader: # ... 前向传播计算损失 loss.backward() optimizer.step() optimizer.zero_grad() # 关键步骤更新EMA影子权重 ema.update() # 验证阶段 ema.apply_shadow() # 应用EMA权重 evaluate(model, val_loader) # 使用平滑后的权重进行评估 ema.restore() # 恢复原始权重继续训练避坑点1clone()的必要性注意update和apply_shadow方法中我们对张量都使用了.clone()。这是至关重要的。在PyTorch中直接赋值self.shadow[name] new_average会导致self.shadow[name]和new_average共享同一块内存。随后new_average被释放或修改会意外地改变影子权重。.clone()创建了一个数据的独立副本确保了影子权重的独立性。避坑点2仅对可训练参数操作我们通过param.requires_grad进行判断。对于固定不变的参数如预训练模型中被冻结的层或者像BatchNorm的running_mean/running_var这种在训练中通过移动平均更新的统计量通常不应对其再做EMA。对统计量做EMA会导致其更新规则被干扰可能影响模型性能。避坑点3初始化的影响在训练初期模型权重变化剧烈且影子权重是从初始权重开始平均的。如果一开始就应用EMA可能会将模型“拉回”到较差的初始点。一个常见的技巧是设置一个预热步数warmup steps在预热期内让decay从一个较小的值如0.9线性或余弦增长到目标值如0.999让EMA在训练初期更快地跟上权重的变化。3.2 方式二使用PyTorch内置的torch.optim.swa_utils官方推荐但需注意细节从PyTorch 1.6开始官方在torch.optim.swa_utils中提供了对随机权重平均SWA的支持而EMA可以看作是SWA的一种特例使用固定衰减系数。AveragedModel类本质上就是一个EMA。import torch.optim.swa_utils as swa_utils model MyModel() # 创建EMA模型avg_fn参数指定了EMA的更新规则 ema_model swa_utils.AveragedModel(model, avg_fnlambda averaged_model_parameter, model_parameter, num_averaged: decay * averaged_model_parameter (1 - decay) * model_parameter) # 训练循环中在optimizer.step()之后 ema_model.update_parameters(model) # 验证时直接使用ema_model with torch.no_grad(): for data, target in val_loader: output ema_model(data) # ... 计算指标优点代码简洁官方维护自动处理设备CPU/GPU和张量拷贝问题。坑点默认情况下AveragedModel会平均模型的所有参数。如果你的模型包含BatchNorm层这可能会出问题。因为BatchNorm层除了权重weight和偏置bias还有在训练中动态更新的running_mean和running_var。对这些统计量做平均是不合适的。解决方案使用swa_utils.get_ema_avg_fn或自定义avg_fn并在创建AveragedModel时传入一个device参数同时更稳妥的做法是在验证前调用swa_utils.update_bn函数基于训练数据重新计算EMA模型的BatchNorm统计量。但这需要额外遍历一遍训练数据增加开销。# 更安全的做法在训练结束后用训练数据更新EMA模型的BN统计量 ema_model swa_utils.AveragedModel(model) # ... 训练循环不断调用 ema_model.update_parameters(model) # 训练结束后 swa_utils.update_bn(train_loader, ema_model, devicedevice) # 更新BN统计量 # 然后再进行最终验证或保存模型个人经验对于研究或快速实验手动实现的ModelEMA类给了我更大的灵活性和可控性特别是对于复杂模型或需要特殊处理如部分参数冻结的情况。而对于生产环境或标准流程使用官方的AveragedModel并妥善处理BN层是更干净、更不易出错的选择。4. EMA超参数调优与效果验证实战设置好EMA代码只是第一步让它真正发挥作用还需要仔细调整超参数并设计实验来验证效果。4.1 核心超参数decay的调优策略decay是EMA唯一的超参数但它对最终模型性能的影响非常显著。我的调优经验遵循以下步骤确定范围对于大多数视觉和NLP任务decay的有效范围通常在[0.99, 0.9999]之间。可以从一个中间值开始比如0.999。结合训练周期decay的选择与总训练迭代次数T密切相关。一个经验法则是希望EMA权重能覆盖足够多的历史迭代。影子权重中最近N ≈ 1 / (1 - decay)次迭代的贡献占主导。例如decay0.99-N≈100次迭代decay0.999-N≈1000次迭代decay0.9999-N≈10000次迭代 你的N应该远小于总迭代次数T否则EMA权重会被早期不成熟的权重过度影响。如果总共只训练5000步用decay0.9999(N10000) 就太大了。网格搜索与观察曲线在一个小的验证集上对几个候选值如0.995, 0.999, 0.9995进行网格搜索。不仅要看最终的验证准确率更要观察验证损失曲线。一个合适的decay应该能使验证损失曲线更平滑、下降更稳定并且最终稳定在一个更低的平台。考虑预热Warmup如前所述在训练初期使用较低的decay或较高的α有助于EMA权重快速跟上模型权重的变化。你可以实现一个动态的decaydef get_current_decay(iter, warmup_iters1000, base_decay0.999): if iter warmup_iters: # 线性从0.9增长到base_decay return 0.9 (base_decay - 0.9) * (iter / warmup_iters) else: return base_decay在每次update时计算当前decay并传入。4.2 效果验证不仅仅是看最终准确率验证EMA是否有效不能只看最终验证集上的一个准确率数字。你需要进行更细致的分析训练vs验证损失曲线对比这是最直观的。在同一张图上绘制使用原始权重和EMA权重计算出的验证损失曲线。一个成功的EMA应该能降低验证损失的波动曲线更平滑。降低验证损失的最终值曲线收敛到更低的点。缓解过拟合训练损失和验证损失之间的差距缩小。权重分布可视化在训练的不同阶段早期、中期、后期分别提取原始权重和EMA权重中某一层的参数如全连接层的权重绘制其直方图或计算其统计量均值、方差。你会发现EMA权重的分布通常更加“集中”方差更小极端值过大或过小的权重更少。这从理论上解释了其更好的泛化性复杂的、过拟合的模型往往拥有一些绝对值非常大的权重EMA平滑了这些极端值。鲁棒性测试对验证集数据施加轻微的扰动如高斯噪声、轻微的模糊、色彩抖动然后分别用原始模型和EMA模型进行测试。EMA模型在扰动下的性能下降通常更小表现出更强的鲁棒性。“快照”集成效应由于EMA权重是历史权重的平均它在某种程度上近似于将训练过程中多个时间点的模型“快照”进行了集成。你可以做一个对比实验单独保存训练过程中几个检查点checkpoint在推理时对它们的预测结果进行平均软投票。这个结果与单一EMA模型的结果进行对比。在很多情况下EMA模型能达到甚至超过多模型集成的效果却只需要存储和运行一个模型效率极高。4.3 一个完整的训练脚本片段示例结合以上所有要点一个整合了动态decay、验证和保存的EMA训练循环核心部分如下import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader import copy class EMA: # ... 使用前面手动实现的ModelEMA类但增加动态decay def __init__(self, model, base_decay0.999, warmup_iters1000): self.model model self.base_decay base_decay self.warmup_iters warmup_iters self.shadow {} self.backup {} for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] param.data.clone().detach() def get_decay(self, iter): if iter self.warmup_iters: return 0.9 (self.base_decay - 0.9) * min(1.0, iter / self.warmup_iters) return self.base_decay def update(self, iter): decay self.get_decay(iter) for name, param in self.model.named_parameters(): if param.requires_grad: self.shadow[name].data.copy_(decay * self.shadow[name].data (1.0 - decay) * param.data) # ... apply_shadow, restore 方法同上 # 初始化 model MyModel().cuda() ema EMA(model, base_decay0.999, warmup_iters2000) optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() best_val_acc 0.0 global_iter 0 for epoch in range(100): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target data.cuda(), target.cuda() optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() # 更新EMA ema.update(global_iter) global_iter 1 # 验证阶段 model.eval() # 使用原始模型验证 orig_val_acc evaluate(model, val_loader) # 使用EMA模型验证 ema.apply_shadow() ema_val_acc evaluate(model, val_loader) # 此时model的权重已被替换为EMA权重 ema.restore() print(fEpoch {epoch}: Orig Val Acc: {orig_val_acc:.4f}, EMA Val Acc: {ema_val_acc:.4f}) # 保存最佳EMA模型 if ema_val_acc best_val_acc: best_val_acc ema_val_acc ema.apply_shadow() best_model_state copy.deepcopy(model.state_dict()) # 保存的是EMA权重 ema.restore() torch.save({ epoch: epoch, model_state_dict: best_model_state, ema_shadow: ema.shadow, # 也可以选择保存影子字典 optimizer_state_dict: optimizer.state_dict(), best_acc: best_val_acc, }, best_ema_model.pth)这个流程清晰地展示了EMA如何无缝嵌入训练循环并通过对比验证准确率来体现其价值。保存模型时务必注意你保存的是应用了影子权重后的模型状态best_model_state还是影子字典本身。前者加载后可直接用于推理后者则需要先初始化EMA类再加载。5. 进阶话题EMA的变体、局限性与相关技术掌握了基础用法后我们可以看看EMA的一些高级变体、它不适用的场景以及与其相关的其他技术。5.1 EMA的常见变体带偏置校正的EMA在训练的最初几步由于影子权重从零或初始值开始其值会偏向于初始值存在偏差。尤其是在decay接近1时这个偏差在早期会很明显。偏置校正通过在早期对EMA值进行缩放来消除这个影响。公式为θ_t_corrected θ_t / (1 - decay^t)。这在Adam等优化器中很常见但在模型权重的EMA中由于我们通常关心训练稳定后的最终权重且预热策略也能缓解此问题所以不常使用。周期性EMAStochastic Weight Averaging, SWASWA可以看作是EMA的一个“节拍器”变体。它不是在每一步都更新影子权重而是以固定的周期如每几个epoch将当前权重加入到平均池中。更新公式是简单的算术平均θ_new (θ_old * n θ_current) / (n 1)。SWA的理论基础是SGD优化路径会在最优解周围的多模盆地中游走平均这些权重可以落在更中心的泛化更好的区域。SWA通常不需要调整decay这样的超参数且在许多任务上表现比固定decay的EMA更鲁棒。PyTorch的swa_utils主要就是为SWA设计的。EMA与学习率调度的协同当使用学习率衰减策略时如StepLR、CosineAnnealing模型权重在后期更新幅度变小。此时EMA的平滑效应会更强。一个有趣的实践是在训练末期可以增大decay例如从0.999增加到0.9999让影子权重变化更慢进一步平滑最后阶段的微小波动有助于模型收敛到更平坦的极小值这通常与更好的泛化性相关。5.2 EMA的局限性什么时候可能没用甚至有害EMA不是银弹在以下场景需要谨慎使用或避免使用训练数据极度干净过拟合风险极低如果你的模型很简单或者数据量极大模型本身就不容易过拟合那么EMA带来的提升可能微乎其微白增加了复杂性。优化器本身已具备强平滑性如果你使用的是像AdamW这样自适应学习率且自带动量对梯度一阶矩估计的优化器它已经在一定程度上平滑了更新过程。再加上EMA效果可能叠加也可能过度平滑导致模型收敛变慢。需要实验验证。对BatchNorm层处理不当这是最大的坑。如前所述对BN层的running_mean/var做EMA会破坏其统计特性。标准做法是EMA只应用于模型的可学习参数权重和偏置而不应用于BN的统计量。在验证/测试时应使用EMA权重但BN层应使用其自身在训练中累积的running_mean/var即原始模型BN层的统计量或者最好在训练结束后用EMA权重重新前向传播一遍训练数据来更新BN统计量update_bn。动态网络结构或权重共享对于结构在训练中发生变化的网络如某些NAS方法或权重共享的模块EMA的更新逻辑可能变得复杂需要特别设计。训练初期在模型权重快速变化的初期阶段过早应用强力的EMA高decay会拖慢影子权重的更新可能不利于模型快速找到有希望的区域。这就是为什么预热策略很重要。5.3 与EMA相关的热词EMA注意力机制最近在搜索EMA时常会看到“EMA注意力机制”这个词。这里需要做一个重要的区分本文讨论的EMA指数移动平均是一种应用于模型参数的平滑技术用于提升模型泛化能力。它是一个训练技巧。EMA注意力机制Efficient Multi-scale Attention这是一种网络模块结构的设计通常出现在计算机视觉的骨干网络中如EMANet Efficient Multi-scale Attention Network。它通过引入多尺度上下文信息和高效的注意力计算来提升模型的特征提取能力。这里的“EMA”是模块的名称缩写其内部可能使用了移动平均的思想来进行特征融合但和本文所述的参数平滑技术是完全不同的概念和应用层面。不要混淆两者。当你看到一篇论文或代码中提到“EMA模块”时需要根据上下文判断它指的是参数平滑技巧还是一种特定的神经网络层。6. 总结与个人心得回顾整个探索过程EMA给我的最大启示是在追求模型性能的路上我们不仅要关注“前沿”和“复杂”更要重视那些被验证有效的“基础”与“简单”。EMA几乎没有增加任何计算开销只是多存储一套权重和一些简单的标量运算却能稳定地带来1-2个百分点的泛化性能提升这在很多竞赛和实际项目中可能就是决定性的优势。从我个人的实战经验来看以下几点心得或许对你有帮助把它变成默认选项对于大多数监督学习任务尤其是视觉和NLP任务我现在会习惯性地在训练脚本里加上EMA。它的收益风险比极高。先跑基线再加EMA在调试新模型或新任务时我通常会先不用EMA跑一个基线观察训练和验证曲线的正常形态。然后再加入EMA对比曲线变化这样能更清晰地看到EMA的效果并帮助我设置合适的decay和预热策略。保存与加载的细节决定成败多少次因为保存和加载EMA模型的状态不对而debug到深夜。务必明确你保存的是什么原始权重、影子字典还是应用了影子的模型状态并在加载时进行对应的恢复操作。写一个清晰的文档字符串或注释说明保存格式至关重要。不要神话它EMA是优秀的正则化工具但它不能替代良好的数据、合适的模型结构和正确的优化器选择。它是“锦上添花”而非“雪中送炭”。如果模型在基础设定下都无法收敛先别指望EMA能拯救世界。最后技术总是在发展。EMA和SWA这类权重平均技术本质上是在损失函数的权重空间里寻找更平坦、泛化更好的解。这与当前关于“平坦极小值”与泛化性的理论研究是相呼应的。理解其背后的思想比单纯调用一个API更有价值。下次当你训练模型时不妨花几分钟加上这几行代码它可能会给你带来意想不到的回报。