PyTorch中EMA需在训练循环中对requires_grad=True的参数及BN缓冲区持续做指数加权平均，影子参数须注册为nn.Module子模块以保证checkpoint保存与恢复，多卡下仅rank0更新，推理用上下文管理器临时替换，torch.compile需禁用EMA更新以避免被跳过。
PyTorch中EMA权重更新的正确实现方式
直接用 torch.nn.Module.load_state_dict() 覆盖原模型参数是错的——EMA本质是指数加权移动平均，必须在训练循环中持续累积，不能靠“保存再加载”模拟。
核心逻辑：对每个可学习参数（param）维护一个影子副本（shadow_param），每次更新按公式 shadow_param = decay * shadow_param + (1 - decay) * param 迭代。注意不是梯度更新，而是参数值本身的滑动平均。
-
decay通常设为 0.999 或 0.9999；值越大，历史权重占比越高，响应新参数越慢 - 只对
requires_grad=True的参数做 EMA，bn.running_mean等缓冲区也要单独处理（否则 BN 层推理会出错） - 初始化影子参数必须和原始参数同设备、同 dtype，否则后续 in-place 操作报
RuntimeError: expected same device
避免常见的EMA状态丢失问题
很多人把 EMA 实例定义在训练函数内部，或没把它加入 torch.nn.Module 子模块，导致 checkpoint 保存时漏掉影子参数，恢复后 EMA 彻底失效。
正确做法是将 EMA 封装成一个继承自 torch.nn.Module 的类，并在主模型中注册为子模块（self.ema = EMAModule(model)）。这样 model.state_dict() 自动包含影子参数，torch.load() 也能一并恢复。
立即学习“Python免费学习笔记（深入）”；
- 不要用普通 Python 字典存影子参数——无法被 PyTorch 的序列化机制识别
- 在
forward()中调用 EMA 会干扰训练流程；EMA 只应在optimizer.step()后、zero_grad()前执行 - 多卡 DDP 训练时，EMA 必须在 rank 0 上更新，其他进程跳过，否则各卡维护不同影子参数
推理时切换到EMA权重的可靠方法
不能简单地把影子参数拷回模型——这会污染原始模型的训练状态。更安全的做法是临时替换参数，用完立刻还原，或直接构建一个只读的 EMA 模型副本。
推荐使用上下文管理器模式：
with ema.apply_shadow():
output = model(x)
其中 apply_shadow() 把影子参数复制到模型，__exit__ 再拷回去。这样既保证推理用的是平滑权重，又不破坏训练连续性。
- 若跳过还原步骤，下次训练 step 会基于被污染的参数计算 EMA，导致偏差累积
- 对
torch.jit.trace或 ONNX 导出，必须先调用ema.copy_to(model)固化权重，再导出，否则导出的是原始模型 - BN 层的
running_mean/running_var缓冲区也要同步 EMA，否则即使权重平滑，推理输出仍抖动
PyTorch 2.0+ 中 torch.compile 与EMA的兼容性
torch.compile 会内联函数并优化参数访问路径，而 EMA 的影子参数常通过 named_parameters() 动态遍历更新——这种反射式操作会被编译器视为副作用，导致 EMA 更新被跳过或执行顺序错乱。
目前稳定方案是：对 EMA 更新逻辑禁用编译，或改用静态参数列表（预先缓存所有需 EMA 的参数引用）。
- 在
torch.compile(model)前，确保 EMA 更新不在被编译的 callable 内；把它放在编译范围外的训练循环里 - 避免在
@torch.compile函数中调用ema.update()，否则可能静默失效，且无 warning 提示 - 验证是否生效：打印某层权重的 id，确认训练后该 id 对应的 tensor 值确实在缓慢趋近于影子值
EMA 的关键不在代码长短，而在参数生命周期的精确控制——影子参数何时初始化、何时更新、何时生效、何时持久化，每一步错位都会让平滑效果归零。尤其要注意 BN 缓冲区和多卡同步这两个最常被忽略的环节。
