机器学习模型性能提升技巧：指数加权平均（EMA）

主要内容

什么是EMA?
为什么EMA在测试过程中使用通常能提升模型表现？
Tensorflow实现
PyTorch实现
Refercences

什么是EMA?

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

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

假设我们得到一个参数θ\thetaθ在不同的 epoch 下的值
θ1,θ2,...,θt\theta_1,\theta_2,...,\theta_tθ1,θ2,...,θt

当训练结束的θ\thetaθ的MovingAverage 就是:

vt=β∗vt−1+(1−β)∗vtv_t=\beta*v_{t-1}+(1-\beta)*v_tvt=β∗vt−1+(1−β)∗vt

β\betaβ代表衰减率，该衰减率用于控制模型更新的速度。

Andrew Ng在Course 2 Improving Deep Neural Networks中讲到，ttt时刻变量vvv的滑动平均值大致等于过去1/(1−β)1/(1−\beta)1/(1−β)个时刻vvv值的平均。

图1 不同 β\betaβ 值做EMA的效果对比（天气预报数据）

当β\betaβ越大时，滑动平均得到的值越和vvv的历史值相关。如果β=0.9\beta=0.9β=0.9，则大致等于过去10个vvv值的平均；如果β=0.99\beta=0.99β=0.99，则大致等于过去100个vvv值的平均。(数学证明先省略，因为作者暂时没理解证明过程==)

滑动平均的好处：
　　
占内存少，不需要保存过去10个或者100个历史vvv值，就能够估计其均值。（当然，滑动平均不如将历史值全保存下来计算均值准确，但后者占用更多内存和计算成本更高）

为什么EMA在测试过程中使用通常能提升模型表现？

滑动平均可以使模型在测试数据上更健壮（robust）。“采用随机梯度下降算法训练神经网络时，使用滑动平均在很多应用中都可以在一定程度上提高最终模型在测试数据上的表现。”

对神经网络边的权重 weights 使用滑动平均，得到对应的影子变量shadow_weights。在训练过程仍然使用原来不带滑动平均的权重 weights，以得到 weights 下一步更新的值，进而求下一步 weights 的影子变量 shadow_weights。之后在测试过程中使用shadow_weights 来代替 weights 作为神经网络边的权重，这样在测试数据上效果更好。因为 shadow_weights 的更新更加平滑，对于：

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

举例来说，设decay=0.999decay=0.999，直观理解，在最后的1000次训练过程中，模型早已经训练完成，正处于抖动阶段，而滑动平均相当于将最后的1000次抖动进行了平均，这样得到的权重会更加robust。

Tensorflow实现

TensorFlow 提供了 tf.train.ExponentialMovingAverage来实现滑动平均。

Example usage when creating a training model:

# Create variables.
var0 = tf.Variable(...)
var1 = tf.Variable(...)
# ... use the variables to build a training model...
...
# Create an op that applies the optimizer.  This is what we usually
# would use as a training op.
opt_op = opt.minimize(my_loss, [var0, var1])# Create an ExponentialMovingAverage object
ema = tf.train.ExponentialMovingAverage(decay=0.9999)with tf.control_dependencies([opt_op]):# Create the shadow variables, and add ops to maintain moving averages# of var0 and var1. This also creates an op that will update the moving# averages after each training step.  This is what we will use in place# of the usual training op.training_op = ema.apply([var0, var1])...train the model by running training_op...

There are two ways to use the moving averages for evaluations:

Build a model that uses the shadow variables instead of the variables.
For this, use the average() method which returns the shadow variable
for a given variable.
Build a model normally but load the checkpoint files to evaluate by using
the shadow variable names. For this use the average_name() method. See
the tf.train.Saver for more
information on restoring saved variables.

Example of restoring the shadow variable values:

# Create a Saver that loads variables from their saved shadow values.
shadow_var0_name = ema.average_name(var0)
shadow_var1_name = ema.average_name(var1)
saver = tf.train.Saver({shadow_var0_name: var0, shadow_var1_name: var1})
saver.restore(...checkpoint filename...)
# var0 and var1 now hold the moving average values

PyTorch实现

PyTorch官方目前没有提供EMA的实现，不过自己实现也不会太复杂，下面提供一个网上大神的实现方法：

class EMA():def __init__(self, decay):self.decay = decayself.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.shadownew_average = (1.0 - self.decay) * x + self.decay * self.shadow[name]self.shadow[name] = new_average.clone()

使用方法，分为初始化、注册和更新三个步骤。

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

Refercences

[1]. 理解滑动平均(exponential moving average)

[2]. EMA 指数滑动平均原理和实现 (PyTorch)
[3]. tf.train.ExponentialMovingAverage