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

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

时间：2019-11-15 20:38:21

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 torchimport torch.nn as nndef 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步骤

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

2.2代码

import torchimport torch.nn as nnfrom torch.utils.data import DataLoaderclass EMA:def __init__(self, model: nn.Module,decay: float = 0.999):self.model = modelself.decay = decayself.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.shadownew_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.shadowself.backup[name] = param.dataparam.data = self.shadow[name]def restore(self):"""恢复模型权重"""for name, param in self.model.named_parameters():if param.requires_grad:assert name in self.backupparam.data = self.backup[name]self.backup = {}