pytorch不提供官方ema模块，直接用state_dict赋值会破坏梯度计算图、导致backward失败；正确做法是在torch.no_grad()中遍历parameters和buffers做原地加权更新。
PyTorch 本身不提供开箱即用的 EMA 模块，但用 torch.no_grad() 和原地更新就能稳定、高效地实现——关键不是“能不能”，而是“怎么避免参数覆盖和梯度污染”。
为什么不能直接用 model.state_dict() 赋值做 EMA？
直接拷贝 state_dict 后做加权平均会丢失参数的梯度计算图引用，且无法保证与模型当前参数内存地址一致；更严重的是，如果在训练循环中频繁调用 load_state_dict()，会破坏 autograd 的张量依赖链，导致后续 backward 失败或梯度为 0。
正确做法是遍历模型参数（model.parameters()）和缓冲区（model.buffers()），对每个 Tensor 做原地更新：
- EMA 更新必须在
torch.no_grad()上下文中进行 - 只对
requires_grad=True的参数做 EMA（通常就是所有可训练参数） - 缓冲区（如 BatchNorm 的
running_mean）也应参与 EMA，否则推理时统计量不准
torch.ema_update() 在 PyTorch 2.0+ 中可用吗？
不可用。截至 PyTorch 2.3，官方仍未提供 torch.ema_update() 或类似 API。社区常见误传源于混淆了内部测试函数或第三方库（如 timm 的 ModelEma）。实际项目中需自行实现，或依赖 torch.optim.swa_utils.AveragedModel（但它是 SWA，非指数衰减）。
一个轻量级实现示例：
class ModelEma:
def __init__(self, model, decay=0.9999):
self.model = model
self.decay = decay
self.shadow = {}
for name, param in model.named_parameters():
if param.requires_grad:
self.shadow[name] = param.data.clone().detach()
<pre class="brush:python;toolbar:false;">@torch.no_grad()
def update(self):
for name, param in self.model.named_parameters():
if param.requires_grad:
self.shadow[name].mul_(self.decay).add_(param.data, alpha=1 - self.decay)
EMA 更新时机：train step 还是 eval step？
必须在每次 optimizer.step() 之后、model.train() 状态下更新。原因有二：
- EMA 本质是对「优化器已接受的梯度更新结果」做平滑，早于
step()会导致滞后一帧 - 若放在
model.eval()后（如验证阶段），此时 BN 层的running_mean/std已冻结，EMA 无法同步这些缓冲区变化 - 注意：EMA 参数不应参与反向传播，所以更新必须在
torch.no_grad()内，且不能出现在 loss 计算路径中
保存与加载 EMA 权重时容易漏掉什么？
保存时只存 shadow 字典或重建后的模型参数，但加载时若直接 load_state_dict() 到原始模型，会覆盖原始参数——这没问题；但若想继续训练 EMA 模型，必须同时恢复 decay 值和 shadow 状态。
更稳妥的做法是封装成独立模型：
def ema_to_model(model, ema_helper):
for name, param in model.named_parameters():
if name in ema_helper.shadow:
param.data.copy_(ema_helper.shadow[name])
真正容易被忽略的是：EMA 的衰减率 decay 本身不是模型参数，但它决定了历史权重占比。不同训练长度下，固定 decay=0.9999 在短训（decay = 1 - (1 - base_decay) * (1 - step / total_steps)。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
