【提分trick】SWA（随机权重平均）和EMA（指数移动平均）

1. SWA随机权重平均

 1.1步骤

1.2代码

2.EMA指数移动平均

2.1步骤

2.2代码

 3.总结

---

在kaggle比赛中，不管是目标检测任务、语义分割任务中，经常能看到SWA（Stochastic Weight Averaging）和EMA（Exponential Moving Average）的身影，今天就来一起学习下。

1. SWA随机权重平均

SWA随机权重平均：在优化的末期取k个优化轨迹上的checkpoints，平均他们的权重，得到最终的网络权重，这样就会使得最终的权重位于flat曲面更中心的位置，缓解权重震荡问题，获得一个更加平滑的解，相比于传统训练有更泛化的解。

​

 1.1步骤

1.给定超参数：

• 循环周期c，代表训练c步就使用SWA进行一次权重平均
• 学习率 α1,α2 ，即周期学习率的上界和下界，论文的实验使用的周期性学习率如下图

2.然后，按照正常的SGD标准流程进行训练，每训练c步，就平均一次权重

3.最后，使用平均的权重 wSWA 权重进行推理。

1.2代码

import torch
import torch.nn as nn


def apply_swa(model: nn.Module,
              checkpoint_list: list,
              weight_list: list,
              strict: bool = True):
    """

    :param model:
    :param checkpoint_list: 要进行swa的模型路径列表
    :param weight_list: 每个模型对应的权重
    :param strict: 输入模型权重与checkpoint是否需要完全匹配
    :return:
    """

    checkpoint_tensor_list = [torch.load(f, map_location='cpu') for f in checkpoint_list]

    for name, param in model.named_parameters():
        try:
            param.data = sum([ckpt['model'][name] * w for ckpt, w in zip(checkpoint_tensor_list, weight_list)])
        except KeyError:
            if strict:
                raise KeyError(f"Can't match '{name}' from checkpoint")
            else:
                print(f"Can't match '{name}' from checkpoint")

    return model

2.EMA指数移动平均

EMA指数移动平均：shadow权重是通过历史的模型权重指数加权平均数来累积的，每次shadow权重的更新都会受上一次shadow权重的影响，所以shadow权重的更新都会带有前几次模型权重的惯性，历史权重越久远，其重要性就越小，这样可以使得权重更新更加平滑。

​

•  为EMA权重，也被称为影子权重；
• 为衰退率，一般为0.999或0.9999；
• 为模型权重

从上述公式来看，shadow权重的更新大部分由累积的权重决定，小部分由当前权重决定。 

2.1步骤

1. 创建EMA平滑的shadow权重（对应EMA对象初始化和register方法）
2. 按照正常的训练流程，反向传播更新模型权重
3. 更新模型权重之后，再执行EMA平滑，更新shadow权重（对应update方法）
4. 重复2-3步，直到valid阶段
5. 备份模型权重，加载shadow权重，使用shadow权重进行模型的valid工作（对应apply_shadow方法）
6. 使用shadow权重作为模型权重，保存模型
7. 恢复模型权重（对应restore方法），继续重复以上步骤2-7。

2.2代码

import torch
import torch.nn as nn
from torch.utils.data import DataLoader


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

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

    def update(self):
        """EMA平滑操作，更新shadow权重"""
        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):
        """使用shadow权重作为模型权重，并创建原模型权重备份"""
        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 = {}

 3.总结

1. EMA需要在每步训练时，同步更新shadow权重，但其计算量与模型的反向传播相比，成本很小，因此实际上并不会拖慢很对模型的训练进度；
2. SWA可以在训练结束，进行手动加权，完全不增加额外的训练成本；
3. 实际使用两者可以配合使用，可以带来一点模型性能提升。

整理不易，欢迎一键三连！！！