再也不用担心过拟合的问题了.docxVIP

SAM的原理是什么？

在梯度下降或任何其他优化算法中，我们的目标是找到一个具有低损失值的参数。但是，与其他常规的优化方法相比，SAM实现了更好的泛化，它将重点放在领域内寻找具有均匀的低损失值的参数(而不是只有参数本身具有低损失值)上。

由于计算邻域参数而不是计算单个参数，损失超平面比其他优化方法更平坦，这反过来增强了模型的泛化。

(左))用SGD训练的ResNet收敛到的一个尖锐的最小值。(右)用SAM训练的相同的ResNet收敛到的一个平坦的最小值。

注意：SAM不是一个新的优化器，它与其他常见的优化器一起使用，比如SGD/Adam。

在Pytorch中实现SAM

在Pytorch中实现SAM非常简单和直接

import?torch

class?SAM(torch.optim.Optimizer):

????def?__init__(self,?params,?base_optimizer,?rho=0.05,?**kwargs):

????????assert?rho?=?0.0,?fInvalidrho,?shouldbenon-negative:?{rho}

????????defaults?=?dict(rho=rho,?**kwargs)

????????super(SAM,?self).__init__(params,?defaults)

????????self.base_optimizer?=?base_optimizer(self.param_groups,?**kwargs)

????????self.param_groups?=?self.base_optimizer.param_groups

????@torch.no_grad()

????def?first_step(self,?zero_grad=False):

????????grad_norm?=?self._grad_norm()

????????for?group?in?self.param_groups:

????????????scale?=?group[rho]?/?(grad_norm?+?1e-12)

????????????for?p?in?group[params]:

????????????????if?p.grad?is?None:?continue

????????????????e_w?=?p.grad?*?scale.to(p)

????????????????p.add_(e_w)??#?climbtothelocalmaximum?w?+?e(w)

????????????????self.state[p][e_w]?=?e_w

????????if?zero_grad:?self.zero_grad()

????@torch.no_grad()

????def?second_step(self,?zero_grad=False):

????????for?group?in?self.param_groups:

????????????for?p?in?group[params]:

????????????????if?p.grad?is?None:?continue

????????????????p.sub_(self.state[p][e_w])??#?getbackto?w?from?w?+?e(w)

????????self.base_optimizer.step()??#?dotheactual?sharpness-aware?update

????????if?zero_grad:?self.zero_grad()

????def?_grad_norm(self):

????????shared_device?=?self.param_groups[0][params][0].device??#?puteverythingonthesamedevice,?incaseofmodelparallelism

????????norm?=?torch.norm(

????????????????????torch.stack([

????????????????????????p.grad.norm(p=2).to(shared_device)

????????????????????????for?group?in?self.param_groups?for?p?in?group[params]

????????????????????????if?p.grad?is?not?None

????????????????????]),

???