SWA(随机权重平均)
[Averaging Weights Leads to Wider Optima and Better Generalization](Averaging Weights Leads to Wider Optima and Better Generalization)
随机权重平均：在优化的末期取k个优化轨迹上的checkpoints，平均他们的权重，得到最终的网络权重，这样就会使得最终的权重位于flat曲面更中心的位置，缓解权重震荡问题，获得一个更加平滑的解，相比于传统训练有更泛化的解。
效果如下：
SWA和EMA
在EMA指数滑动平均（Exponential Moving Average）我们讨论了指数滑动平均，可以发现SWA和EMA是有相似之处：
- 都是在训练之外的操作，不影响训练过程。
- 与集成学习类似，都是一种权值的平均，EMA是一种指数平均，会赋予近期更多的权重，SWA则是平均赋权重。
所以这里参考了的SWA实现，添加了EMA的实现，两者不同在于影子权值的更新方式。
class WeightAverage(Optimizer):
def __init__(self, optimizer, wa_start=None, wa_freq=None, wa_lr=None, mode='swa'):
"""实现参考：https://github.com/pytorch/contrib/blob/master/torchcontrib/optim/swa.py
论文：Averaging Weights Leads to Wider Optima and Better Generalization
两种权重平均的方式 swa 和 ema
两种模式：自动模式和手动模式
参数:
optimizer (torch.optim.Optimizer): optimizer to use with SWA
wa_start (int): SWA开始应用的step
wa_freq (int): 更新SWA的频数
wa_lr (float): 自动模式：从swa_start开始应用
"""
if isinstance(mode, float):
self.mode = 'ema'
self.beta = mode
else:
self.mode = mode
self._auto_mode, (self.wa_start, self.wa_freq) = self._check_params(wa_start, wa_freq)
self.wa_lr = wa_lr
# 参数检查
if self._auto_mode:
if wa_start < 0:
raise ValueError("Invalid wa_start: {}".format(wa_start))
if wa_freq < 1:
raise ValueError("Invalid wa_freq: {}".format(wa_freq))
else:
if self.wa_lr is not None:
warnings.warn("Some of wa_start, wa_freq is None, ignoring wa_lr")
self.wa_lr = None
self.wa_start = None
self.wa_freq = None
if self.wa_lr is not None and self.wa_lr < 0:
raise ValueError("Invalid WA learning rate: {}".format(wa_lr))
self.optimizer = optimizer
self.defaults = self.optimizer.defaults
self.param_groups = self.optimizer.param_groups
self.state = defaultdict(dict)
self.opt_state = self.optimizer.state
for group in self.param_groups:
# ema 不需要保存已经平均的个数,为了兼容swa不修改
group['n_avg'] = 0
group['step_counter'] = 0
@staticmethod
def _check_params(swa_start, swa_freq):
"""检查参数，确认执行模式，并将参数转为int
"""
params = [swa_start, swa_freq]
params_none = [param is None for param in params]
if not all(params_none) and any(params_none):
warnings.warn("Some of swa_start, swa_freq is None, ignoring other")
for i, param in enumerate(params):
if param is not None and not isinstance(param, int):
params[i] = int(param)
warnings.warn("Casting swa_start, swa_freq to int")
return not any(params_none), params
def _reset_lr_to_swa(self):
"""应用wa学习率
"""
if self.wa_lr is None:
return
for param_group in self.param_groups:
if param_group['step_counter'] >= self.wa_start:
param_group['lr'] = self.wa_lr
def update_swa_group(self, group):
"""更新一组参数的wa: 随机权重平均或者指数滑动平均
"""
for p in group['params']:
param_state = self.state[p]
if 'wa_buffer' not in param_state:
param_state['wa_buffer'] = torch.zeros_like(p.data)
buf = param_state['wa_buffer']
if self.mode == 'swa':
virtual_decay = 1 / float(group["n_avg"] + 1)
diff = (p.data - buf) * virtual_decay # buf + (p-buf) / (n+1) = (p + n*buf) / (n+1)
buf.add_(diff)
else:
buf.mul_(self.beta).add_((1-self.beta) * p.data)
group["n_avg"] += 1
def update_swa(self):
"""手动模式：更新所有参数的swa
"""
for group in self.param_groups:
self.update_swa_group(group)
def swap_swa_sgd(self):
"""1.交换swa和模型的参数 2.训练结束时和评估时调用
"""
for group in self.param_groups:
for p in group['params']:
param_state = self.state[p]
if 'wa_buffer' not in param_state:
warnings.warn("WA wasn't applied to param {}; skipping it".format(p))
continue
buf = param_state['wa_buffer']
tmp = torch.empty_like(p.data)
tmp.copy_(p.data)
p.data.copy_(buf)
buf.copy_(tmp)
def step(self, closure=None):
"""1.梯度更新 2.如果是自动模式更新swa参数
"""
self._reset_lr_to_swa()
loss = self.optimizer.step(closure)
for group in self.param_groups:
group["step_counter"] += 1
steps = group["step_counter"]
if self._auto_mode:
if steps > self.wa_start and steps % self.wa_freq == 0:
self.update_swa_group(group)
return loss
def state_dict(self):
"""打包 opt_state 优化器状态，swa_state SWA状态，param_groups 参数组
"""
opt_state_dict = self.optimizer.state_dict()
wa_state = {(id(k) if isinstance(k, torch.Tensor) else k): v
for k, v in self.state.items()}
opt_state = opt_state_dict["state"]
param_groups = opt_state_dict["param_groups"]
return {"opt_state": opt_state, "wa_state": wa_state,
"param_groups": param_groups}
def load_state_dict(self, state_dict):
"""加载swa和优化器的状态参数
"""
wa_state_dict = {"state": state_dict["wa_state"],
"param_groups": state_dict["param_groups"]}
opt_state_dict = {"state": state_dict["opt_state"],
"param_groups": state_dict["param_groups"]}
super(WeightAverage, self).load_state_dict(wa_state_dict)
self.optimizer.load_state_dict(opt_state_dict)
self.opt_state = self.optimizer.state
def add_param_group(self, param_group):
"""将一组参数添加到优化器的 `param_groups`.
"""
param_group['n_avg'] = 0
param_group['step_counter'] = 0
self.optimizer.add_param_group(param_group)
@staticmethod
def bn_update(loader, model, device=None):
"""更新 BatchNorm running_mean, running_var
"""
if not _check_bn(model):
return
was_training = model.training
model.train()
momenta = {}
model.apply(_reset_bn)
model.apply(lambda module: _get_momenta(module, momenta))
n = 0
for input in loader:
if isinstance(input, (list, tuple)):
input = input[0]
b = input.size(0) # batch_size
momentum = b / float(n + b)
for module in momenta.keys():
module.momentum = momentum
if device is not None:
input = input.to(device)
model(input)
n += b
model.apply(lambda module: _set_momenta(module, momenta))
model.train(was_training)
# BatchNorm utils
def _check_bn_apply(module, flag):
if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
flag[0] = True
def _check_bn(model):
flag = [False]
model.apply(lambda module: _check_bn_apply(module, flag))
return flag[0]
def _reset_bn(module):
if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
module.running_mean = torch.zeros_like(module.running_mean)
module.running_var = torch.ones_like(module.running_var)
def _get_momenta(module, momenta):
if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
momenta[module] = module.momentum
def _set_momenta(module, momenta):
if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
module.momentum = momenta[module]
Stochastic Weight Averaging in PyTorch