深度学习 性能提升技巧--指数加权平均（EMA）Pytorch实现

• 什么是EMA（ExponentialMovingAverage）

在采用 SGD 或者其他的一些优化算法 (Adam, Momentum) 训练神经网络时，通常会使用一个叫 ExponentialMovingAverage (EMA) 的方法，中文名叫指数滑动平均。 它的意义在于利用滑动平均的参数来提高模型在测试数据上的健壮性。

滑动平均(exponential moving average)，或者叫做指数加权平均(exponentially weighted moving average)，可以用来估计变量的局部均值，使得变量的更新与一段时间内的历史取值有关。

滑动平均可以看作是变量的过去一段时间取值的均值，相比对变量直接赋值而言，滑动平均得到的值在图像上更加平缓光滑，抖动性更小，不会因为某次的异常取值而使得滑动平均值波动很大

随机梯度下降，更平滑的更新说明不会偏离最优点很远；
梯度下降 batch gradient decent，影子变量作用可能不大，因为梯度下降的方向已经是最优的了，loss 一定减小；
mini-batch gradient decent，可以尝试滑动平均，因为mini-batch gradient decent 对参数的更新也存在抖动。

• EMA定义，Pytorch实现

class EMA():
    def __init__(self, mu):
        self.mu = mu
        self.shadow = {}

    def register(self, name, val):
        self.shadow[name] = val.clone()

    def get(self, name):
        return self.shadow[name]

    def update(self, name, x):
        assert name in self.shadow
        new_average = (1.0 - self.mu) * x + self.mu * self.shadow[name]
        self.shadow[name] = new_average.clone()

 

• 训练过程中怎么用

train时：

1.注册---拿到model的参数，加入

ema = EMA(0.999)
    for name, param in model.named_parameters():
        if param.requires_grad:
            ema.register(name, param.data)

2.更新--喂batch后，模型更新参数之后，取到model的param,执行param的指数加权平均操作

       loss += batch_loss.item()
        batch_loss.backward()
        optimizer.step()
        
        for name, param in model.named_parameters():
            if param.requires_grad:
                ema.update(name, param.data)

3.测试时候要把加权的更新去掉，直接去参数的值

    backup_params = EMA(0)#注意这里衰减参数值为0
    for name, param in model.named_parameters():
        if param.requires_grad:
            backup_params.register(name, param.data)
            param.data.copy_(ema.get(name))

 

 

其他博客

知乎 EMA

EMA