指数移动平均（EMA）策略

文章目录

• • 概述
  • 具体步骤
  • 代码实现

概述

指数移动平均（EMA）是一种加权移动平均的方法，它给予近期数据更高的权重，同时也考虑到了历史数据的影响。在神经网络领域，EMA 常被用于对模型参数进行平滑处理，使得网络模型在训练过程中能够更加稳定且泛化能力可能得到提升。

具体步骤

假设我们有一个神经网络模型，其参数为 θ \theta θ（例如权重矩阵和偏置向量等），我们要使用EMA策略来更新这些参数。

1. 初始化EMA参数：
   设 θ e m a \theta_{ema} θema​ 为EMA后的参数，初始时，令 θ e m a = θ \theta_{ema} = \theta θema​=θ，即EMA后的参数初始值与原始模型参数相同。确定一个衰减率（decay rate），通常记为 α \alpha α，取值范围一般在0到1之间，常见的值如0.99、0.999等。这个衰减率决定了历史数据在EMA计算中的权重占比， α \alpha α 越接近1，给予历史数据的权重就越高。
2. 在每次训练迭代中更新EMA参数：
   在每次训练迭代（例如，每处理一个批次的数据）后，我们按照以下公式更新EMA参数： θ e m a = α × θ e m a + ( 1 − α ) × θ \theta_{ema} = \alpha \times \theta_{ema} + (1 - \alpha) \times \theta θema​=α×θema​+(1−α)×θ
   θ \theta θ 是当前训练迭代后通过常规优化算法（如Adam）更新得到的模型参数， θ e m a \theta_{ema} θema​ 是经过EMA策略更新后的参数， α \alpha α为衰减率。

代码实现

import torch
import torch.nn as nn
from collections import OrderedDict

# 定义一个简单的神经网络模型

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

@torch.no_grad()
def update_ema(ema_model, model, decay=0.9999):
   """
    Step the EMA model towards the current model.

    Args:
        ema_model: The EMA model whose parameters will be updated.
        model: The current model from which we take the updated parameters.
        decay: The decay rate for EMA update. A value closer to 1 means slower EMA update.
    """
    ema_params = OrderedDict(ema_model.named_parameters())
    model_params = OrderedDict(model.named_parameters())

    for name, param in model_params.items():
        # Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)


# 创建模型实例
model = SimpleNet()
# 定义优化器，这里以Adam为例
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 将EMA模型的参数初始化为和原始模型一样
ema_model.load_state_dict(model.state_dict())
# 设置EMA的衰减率decay，一般取值在0到1之间，越接近1，EMA更新越缓慢
decay = 0.9999
# 定义字典用于存储模型参数
model_params = {}
# 在训练循环中使用EMA更新参数
for epoch in range(100):
    # 假设这里有训练数据加载和前向传播、计算损失等常规训练步骤
    optimizer.zero_grad()
    inputs = torch.randn(32, 10)
    targets = torch.randn(32, 1)
    outputs = model(inputs)
    loss = nn.MSELoss()(outputs, targets)
    loss.backward()
    optimizer.step()
    update_ema(ema_model, model, decay)
model_parmars['model'] = model.state_dict()
model_parmars['ema_model'] = ema_model.state_dict()
torch.save(model_parmars, 'SimpleNet.pth')

通过指数移动平均（EMA）策略更新网络模型参数，可以在一定程度上提高模型的稳定性和泛化能力，是一种在网络模型训练中较为常用的技巧。