方法说明

指数移动平均方法，基于近期的数据更高权重的平均方法。

对于n个数据[θ1,θ2,…,θn][\theta_1, \theta_2, \dots, \theta_n]

  • 普通的平均数为：vˉ=1n∑i=1nθi\bar{v} = \frac{1}{n}\sum_{i=1}^n\theta_i
  • EMA为：vt=β⋅vt−1+(1−β)⋅θtv_t=\beta\cdot v_{t-1}+(1-\beta)\cdot\theta_t，其中vtv_t代表前tt条的平均值（v0=0v_0=0），β\beta是加权权重值，一般为0.9到0.999，反正就是非常的大。

在深度学习的梯度下降过程中，θt\theta_t是模型在tt时刻的权重，vtv_t是tt时刻的影子权重，这个影子权重不会参加训练，但是会被一直维护，而是用来进行优化。在模型训练的最后阶段，由于它在最优点抖动，所以一般取影子权重的平均就好了。

实现代码

下面是一个简单的实现代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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()