EMA原理和pytorch实现

EMA的原理和pytorch实现

EMA 定义

指数移动平均也叫权重移动平均，是一种给予近期数据更高权重的平均方法。

在深度学习的优化中的EMA

在深度学习的优化过程中，$\theta_t$ 是 t 时刻的模型权重 weight， $\upsilon_t$ 是 t 时刻的影子权重。在梯度下降的过程中，会一直维护着这个影子权重，但是这个影子权重不会参与训练。基本的假设是模型权重在最后 n 步内，会在实际的最优点处抖动，所以我们去最后 n 步的平均，能使模型更加的鲁棒。

EMA为何有效

因为在训练的时候，会使用验证集来衡量模型精度，但是验证集和测试集并不完全一致，在训练后期阶段，模型可能已经在测试集最佳精度附近波动，所以使用ema的结果会比使用单一结果更可靠。

Pytorch实现

class EMA():
    def __init__(self, model, decay):
        self.model = model
        self.decay = decay
        self.shadow = {}
        self.backup = {}

    def register(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def update(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
                self.shadow[name] = new_average.clone()

    def apply_shadow(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                self.backup[name] = param.data
                param.data = self.shadow[name]

    def restore(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

# 初始化
ema = EMA(model, 0.999)
ema.register()

# 训练过程中，更新完参数后，同步update shadow weights
def train():
    optimizer.step()
    ema.update()

# eval前，apply shadow weights；eval之后，恢复原来模型的参数
def evaluate():
    ema.apply_shadow()
    # evaluate
    ema.restore()