ema-pytorch

EMA - Pytorch

一种简单的方法来跟踪你的Pytorch模型的指数移动平均（EMA）版本

安装

$ pip install ema-pytorch

使用方法

import torch
from ema_pytorch import EMA

# 你的神经网络作为一个pytorch模块

net = torch.nn.Linear(512, 512)

# 包装你的神经网络，指定衰减因子（beta）

ema = EMA(
    net,
    beta = 0.9999,              # 指数移动平均因子
    update_after_step = 100,    # 只有在调用.update()这么多次之后才开始更新
    update_every = 10,          # 实际更新的频率，以节省计算资源（每10次.update()调用才更新一次）
)

# 改变你的网络，使用SGD或其他方法

with torch.no_grad():
    net.weight.copy_(torch.randn_like(net.weight))
    net.bias.copy_(torch.randn_like(net.bias))

# 你将在移动平均包装器上调用更新函数

ema.update()

# 然后，之后你可以像使用你的网络一样调用EMA模型

data = torch.randn(1, 512)

output     = net(data)
ema_output = ema(data)

# 如果你想保存你的ema模型，建议你保存整个包装器
# 因为它包含了已经进行的步数（其中有一个预热逻辑，由@crowsonkb推荐，并在多个项目中得到验证）
# 然而，如果你希望访问带有EMA的模型副本，它将位于ema.ema_model

要使用Karras等人在<a href="https://arxiv.org/abs/2312.02696">最近的一篇论文</a>中提出的事后合成EMA，请按照以下示例操作

import torch
from ema_pytorch import PostHocEMA

# 你的神经网络作为一个pytorch模块

net = torch.nn.Linear(512, 512)

# 包装你的神经网络，指定sigma_rels或gammas

emas = PostHocEMA(
    net,
    sigma_rels = (0.05, 0.3),           # 一个包含多个EMA超参数的元组。这里至少需要2个来合成新的EMA
    update_every = 10,                  # 实际更新的频率，以节省计算资源（每10次.update()调用才更新一次）
    checkpoint_every_num_steps = 10,
    checkpoint_folder = './post-hoc-ema-checkpoints'  # 保存每个sigma_rel（gamma）在不同时间步的检查点的文件夹，用于训练后合成新的EMA模型
)

net.train()

for _ in range(1000):
    # 改变你的网络，使用SGD或其他方法

    with torch.no_grad():
        net.weight.copy_(torch.randn_like(net.weight))
        net.bias.copy_(torch.randn_like(net.bias))

    # 你将在移动平均包装器上调用更新函数

    emas.update()

# 现在你有了一些检查点
# 你可以合成一个具有不同sigma_rel的EMA模型（比如0.15）

synthesized_ema = emas.synthesize_ema_model(sigma_rel = 0.15)

# 使用合成的EMA输出

data = torch.randn(1, 512)

synthesized_ema_output = synthesized_ema(data)

引用

@article{Karras2023AnalyzingAI,
    title   = {Analyzing and Improving the Training Dynamics of Diffusion Models},
    author  = {Tero Karras and Miika Aittala and Jaakko Lehtinen and Janne Hellsten and Timo Aila and Samuli Laine},
    journal = {ArXiv},
    year    = {2023},
    volume  = {abs/2312.02696},
    url     = {https://api.semanticscholar.org/CorpusID:265659032}
}