ema-pytorch

收录理由

用过PyTorch的人多半自己写过指数移动平均（EMA）的小工具，但那种临时拼凑的代码总差一口气。这个库把任意nn.Module包一层，就能维护一份影子权重，衰减系数、预热步数和更新频率都可调，训练时旁边挂一个稳定的评估模型或教师模型，省心不少。它还实现了Karras等人提出的后验合成EMA、Switch EMA变体，以及能把EMA输出路由到嵌套子模块的包装器，适合自监督场景。扩散模型训练和任何靠平滑权重提升最终效果的项目都能直接受益，更新循环不用自己操心。与其调试自己那套，不如用这个维护良好的实现，能省下不少时间。