【深度学习实战（25）】搭建训练框架之ModelEMA

一、什么是ModelEMA：

在深度学习中，经常会使用EMA（指数移动平均）这个方法对模型的参数做平均，以求提高测试指标并增加模型鲁棒。
指数移动平均（Exponential Moving Average）也叫权重移动平均（Weighted Moving Average），是一种给予近期数据更高权重的平均方法。

二、如何实现ModelEMA

创建EMA eval mode，去并行化

self.ema = deepcopy(de_parallel(model)).eval() 
1

EMA更新次数

self.updates = updates
1

根据更新次数，获取衰减系数

self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)
1

去掉梯度，ema不需要梯度

for p in self.ema.parameters():
    p.requires_grad_(False)
1
2

EMA更新次数+1

self.updates += 1
1

根据更新次数，获取衰减系数

d = self.decay(self.updates)
1

根据衰减系数，当前模型（去并行化）来修改当前ema模型

msd = de_parallel(model).state_dict()  # model state_dict
            for k, v in self.ema.state_dict().items():
                if v.dtype.is_floating_point:
                    v *= d
                    v += (1 - d) * msd[k].detach()
1
2
3
4
5

三、ModelEMA完整实现

#----------------------#
#   判断是否并行训练模式
#----------------------#
def is_parallel(model):
    # Returns True if model is of type DP or DDP
    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)

#----------------------#
#   去并行训练模式
#----------------------#
def de_parallel(model):
    # De-parallelize a model: returns single-GPU model if model is of type DP or DDP
    return model.module if is_parallel(model) else model

#----------------------#
#   模型拷贝
#----------------------#
def copy_attr(a, b, include=(), exclude=()):
    # Copy attributes from b to a, options to only include [...] and to exclude [...]
    for k, v in b.__dict__.items():
        if (len(include) and k not in include) or k.startswith('_') or k in exclude:
            continue
        else:
            setattr(a, k, v)


class ModelEMA:
    """ Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
    Keeps a moving average of everything in the model state_dict (parameters and buffers)
    For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
    """

    def __init__(self, model, decay=0.9999, tau=2000, updates=0):
        #----------------------#
        #   创建EMA eval mode，去并行化
        #----------------------#
        self.ema = deepcopy(de_parallel(model)).eval()  # FP32 EMA
        #----------------------#
        #   EMA更新次数
        #----------------------#
        self.updates = updates
        #----------------------#
        #   根据更新次数，获取衰减系数
        #----------------------#
        self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)
        #----------------------#
        #   去掉梯度，ema不需要梯度
        #----------------------#
        for p in self.ema.parameters():
            p.requires_grad_(False)

    #----------------------#
    #   根据ema更新次数获取衰减系数，再根据衰减系数和当前模型（去并行化）来修改当前ema模型
    #----------------------#
    def update(self, model):
        # Update EMA parameters
        with torch.no_grad():
            #----------------------#
            #   EMA更新次数+1
            #----------------------#
            self.updates += 1
            #----------------------#
            #   根据更新次数，获取衰减系数
            #----------------------#
            d = self.decay(self.updates)
            print('decay:',d)
            dict_decay.append(d)

            #----------------------#
            #   根据衰减系数，当前模型（去并行化）来修改当前ema模型
            #----------------------#
            msd = de_parallel(model).state_dict()  # model state_dict
            for k, v in self.ema.state_dict().items():
                if v.dtype.is_floating_point:
                    v *= d
                    v += (1 - d) * msd[k].detach()

    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
        # Update EMA attributes
        copy_attr(self.ema, model, include, exclude)
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

四、ModelEMA在训练框架中的使用

#----------------------#
#   搭建训练框架
#----------------------#
model = models.AlexNet()
model = model.train()
#----------------------#
#   创建EMA模型
#----------------------#
ema = ModelEMA(model)


num_train_data = 50
batch_size = 10
epoch_step = num_train_data // batch_size

Init_epoch = 50
Total_epoch = 60

#----------------------#
#   记录EMA更新次数
#----------------------#
ema.updates = Init_epoch * epoch_step

#----------------------#
#   训练
#----------------------#
for epoch in range(Init_epoch, Total_epoch):
    dict_epoch.append(epoch)
    for iter in range(epoch_step):
        #----------------------#
        #   根据ema更新次数获取衰减系数，再根据衰减系数和当前模型（去并行化）来修改当前ema模型
        #----------------------#
        ema.update(model)

#----------------------#
#   验证
#----------------------#
#----------------------#
#   获取EMA eval mode，去并行化
#----------------------#
model = ema.ema
for epoch in range(Init_epoch, Total_epoch):
    for iter in range(epoch_step):
        pass

#----------------------#
#   保存权重
#----------------------#
#----------------------#
#   获取EMA模型的权重
#----------------------#
save_state_dict = ema.ema.state_dict()
path = "yourpath"
torch.save(save_state_dict,path)
print('dene')
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
48
49
50
51
52
53
54
55

五、完整代码

import torch
import math
import torch.nn as nn
from copy import deepcopy
from torchvision import models
import matplotlib.pyplot as plt

dict_decay = []
dict_update_num = []

#----------------------#
#   判断是否并行训练模式
#----------------------#
def is_parallel(model):
    # Returns True if model is of type DP or DDP
    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)

#----------------------#
#   去并行训练模式
#----------------------#
def de_parallel(model):
    # De-parallelize a model: returns single-GPU model if model is of type DP or DDP
    return model.module if is_parallel(model) else model

#----------------------#
#   模型拷贝
#----------------------#
def copy_attr(a, b, include=(), exclude=()):
    # Copy attributes from b to a, options to only include [...] and to exclude [...]
    for k, v in b.__dict__.items():
        if (len(include) and k not in include) or k.startswith('_') or k in exclude:
            continue
        else:
            setattr(a, k, v)


class ModelEMA:
    """ Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
    Keeps a moving average of everything in the model state_dict (parameters and buffers)
    For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
    """

    def __init__(self, model, decay=0.9999, tau=2000, updates=0):
        #----------------------#
        #   创建EMA eval mode，去并行化
        #----------------------#
        self.ema = deepcopy(de_parallel(model)).eval()  # FP32 EMA
        #----------------------#
        #   EMA更新次数
        #----------------------#
        self.updates = updates
        #----------------------#
        #   根据更新次数，获取衰减系数
        #----------------------#
        self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)
        #----------------------#
        #   去掉梯度，ema不需要梯度
        #----------------------#
        for p in self.ema.parameters():
            p.requires_grad_(False)

    #----------------------#
    #   根据ema更新次数获取衰减系数，再根据衰减系数和当前模型（去并行化）来修改当前ema模型
    #----------------------#
    def update(self, model):
        # Update EMA parameters
        with torch.no_grad():
            #----------------------#
            #   EMA更新次数+1
            #----------------------#
            self.updates += 1
            #----------------------#
            #   根据更新次数，获取衰减系数
            #----------------------#
            d = self.decay(self.updates)
            print('decay:',d)
            dict_decay.append(d)

            #----------------------#
            #   根据衰减系数，当前模型（去并行化）来修改当前ema模型
            #----------------------#
            msd = de_parallel(model).state_dict()  # model state_dict
            for k, v in self.ema.state_dict().items():
                if v.dtype.is_floating_point:
                    v *= d
                    v += (1 - d) * msd[k].detach()

    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
        # Update EMA attributes
        copy_attr(self.ema, model, include, exclude)


#----------------------#
#   搭建训练框架
#----------------------#
model = models.AlexNet()
model = model.train()
#----------------------#
#   创建EMA模型
#----------------------#
ema = ModelEMA(model)


num_train_data = 100
batch_size = 10
epoch_step = num_train_data // batch_size

Init_epoch = 50
Total_epoch = 300

#----------------------#
#   记录EMA更新次数
#----------------------#
ema.updates = Init_epoch * epoch_step

#----------------------#
#   训练
#----------------------#
num_update = 0
for epoch in range(Init_epoch, Total_epoch):
    for iter in range(epoch_step):
        #----------------------#
        #   根据ema更新次数获取衰减系数，再根据衰减系数和当前模型（去并行化）来修改当前ema模型
        #----------------------#
        ema.update(model)
        num_update += 1
        dict_update_num.append(num_update)


#----------------------#
#   验证
#----------------------#
#----------------------#
#   获取EMA eval mode，去并行化
#----------------------#
model = ema.ema
for epoch in range(Init_epoch, Total_epoch):
    for iter in range(epoch_step):
        pass

#----------------------#
#   保存权重
#----------------------#
#----------------------#
#   获取EMA模型的权重
#----------------------#
save_state_dict = ema.ema.state_dict()
path = "yourpath"
#torch.save(save_state_dict,path)
print('dene')



# -----------------------------------------------#
#   save EMA decay figure
# -----------------------------------------------#
plt.figure()
plt.title('EMA decay during training')
plt.plot(dict_update_num, dict_decay, label="EMA decay")
plt.legend()
plt.grid()
plt.draw()
plt.savefig('EMA decay')
plt.show()
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164

EMA decay 曲线变化图