（自适应）SAM 优化器

锐度感知最小化：高效提升泛化能力

~ 基于 PyTorch ~

---

SAM 同时最小化损失值和损失的锐度。具体而言，它寻找位于具有均匀低损失的邻域中的参数。SAM 能够提升模型的泛化能力，并在多个数据集上达到 当前最优性能。此外，它对标签噪声的鲁棒性与专门针对带噪声标签学习的当前最优方法相当。

这是一个关于 锐度感知最小化：高效提升泛化能力 和 ASAM：用于深度神经网络尺度不变学习的自适应锐度感知最小化 的 非官方 代码库。从实现上看，SAM 类是一个轻量级封装，用于计算正则化的“锐度感知”梯度，该梯度由底层优化器（如带有动量的 SGD）使用。本仓库还包含一个简单的 CIFAR-10 WRN 示例；作为概念验证，它在该数据集上的表现优于带有动量的 SGD。

_{训练结束时，ResNet 在有 SAM 和无 SAM 情况下的损失景观。锐度感知更新会导致更宽的极小值区域，从而带来更好的泛化性能。}

使用方法

在你的训练流程中使用 SAM 应该非常直观。需要注意的是，训练速度会减慢一倍，因为 SAM 需要两次前向和反向传播来估计“锐度感知”梯度。如果你使用梯度裁剪，请确保只调整梯度的大小，而不改变其方向。

from sam import SAM
...

model = YourModel()
base_optimizer = torch.optim.SGD  # 定义用于“锐度感知”更新的优化器
optimizer = SAM(model.parameters(), base_optimizer, lr=0.1, momentum=0.9)
...

for input, output in data:

  # 第一次前向-反向传播
  loss = loss_function(output, model(input))  # 使用此损失记录任何训练统计信息
  loss.backward()
  optimizer.first_step(zero_grad=True)
  
  # 第二次前向-反向传播
  loss_function(output, model(input)).backward()  # 确保完成完整的前向传播
  optimizer.second_step(zero_grad=True)
...

另一种使用方式：基于闭包的单次 step 函数。这种方式提供了与 PyTorch 原生优化器（如 LBFGS）类似的 API（由 @rmcavoy 友好建议）：

from sam import SAM
...

model = YourModel()
base_optimizer = torch.optim.SGD  # 定义用于“锐度感知”更新的优化器
optimizer = SAM(model.parameters(), base_optimizer, lr=0.1, momentum=0.9)
...

for input, output in data:
  def closure():
    loss = loss_function(output, model(input))
    loss.backward()
    return loss

  loss = loss_function(output, model(input))
  loss.backward()
  optimizer.step(closure)
  optimizer.zero_grad()
...

训练技巧

• @hjq133：如果使用批归一化，上述用法可能会出现问题。运行统计会在两次前向传播中都被计算，但它们应该只在第一次计算。一种可能的解决方案是将 BN 动量设为零（由 @ahmdtaha 友好建议），以避免在第二次传播中更新运行统计。示例用法见 example/train.py 的第 51 和 58 行：

for batch in dataset.train:
  inputs, targets = (b.to(device) for b in batch)

  # 第一次前向-反向步骤
  enable_running_stats(model)  # <- 这是关键行
  predictions = model(inputs)
  loss = smooth_crossentropy(predictions, targets)
  loss.mean().backward()
  optimizer.first_step(zero_grad=True)

  # 第二次前向-反向步骤
  disable_running_stats(model)  # <- 这是关键行
  smooth_crossentropy(model(inputs), targets).mean().backward()
  optimizer.second_step(zero_grad=True)

• @evanatyourservice：如果计划在多 GPU 上训练，论文指出：“为了在多个加速器之间并行化时计算 SAM 更新，我们将每个数据批次均匀地分配到各个加速器，在每个加速器上独立计算 SAM 梯度，然后对得到的子批次 SAM 梯度取平均，以获得最终的 SAM 更新。”这可以通过以下代码实现：

for input, output in data:
  # 第一次前向-反向传播
  loss = loss_function(output, model(input))
  with model.no_sync():  # <- 这是关键行
    loss.backward()
  optimizer.first_step(zero_grad=True)
  
  # 第二次前向-反向传播
  loss_function(output, model(input)).backward()
  optimizer.second_step(zero_grad=True)

• @evanatyourservice：据报道，自适应 SAM 的性能优于原始 SAM。ASAM 论文建议为自适应更新使用更高的 rho 参数（约为原始值的 10 倍）。

• @mlaves：学习率调度应应用于基础优化器，或者你可以使用 SAM 并通过单次 step 调用（结合闭包）来实现：

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer.base_optimizer, T_max=200)

• @AlbertoSabater：与 PyTorch Lightning 集成时，可以这样编写 training_step 函数：

def training_step(self, batch, batch_idx):
    optimizer = self.optimizers()

    # 第一次前向-反向传播
    loss_1 = self.compute_loss(batch)
    self.manual_backward(loss_1, optimizer)
    optimizer.first_step(zero_grad=True)

    # 第二次前向-反向传播
    loss_2 = self.compute_loss(batch)
    self.manual_backward(loss_2, optimizer)
    optimizer.second_step(zero_grad=True)

    return loss_1

文档

SAM.__init__

参数	描述
params (可迭代对象)	需要优化的参数的可迭代对象，或定义参数组的字典
base_optimizer (torch.optim.Optimizer)	执行“尖锐度感知”更新的基础优化器
rho (float, 可选)	用于计算最大损失的邻域大小 (默认: 0.05)
adaptive (bool, 可选)	如果您希望使用基于元素级别的自适应 SAM 的实验性实现，请将此参数设置为 True (默认: False)
**kwargs	传递给 base_optimizer 的 __init__ 方法的关键字参数

SAM.first_step

执行第一步优化，找到局部 rho 邻域内损失最大的权重。

参数	描述
zero_grad (bool, 可选)	如果您希望在此步骤后自动清零所有梯度，请将其设置为 True (默认: False)

SAM.second_step

执行第二步优化，根据损失曲面上（局部）最高点的梯度来更新原始权重。

参数	描述
zero_grad (bool, 可选)	如果您希望在此步骤后自动清零所有梯度，请将其设置为 True (默认: False)

SAM.step

在一次调用中同时执行两个优化步骤。此函数是显式调用 SAM.first_step 和 SAM.second_step 的替代方案。

参数	描述
closure (可调用对象)	闭包应针对优化后的模型执行一次完整的前向和反向传播 (默认: None)

实验

我已验证 SAM 在 CIFAR10 数据集上运行的简单 WRN 16-8 模型上有效；您可以通过运行 train.py 来复现该实验。Wide-ResNet 仅通过标签平滑和最基本的图像增强（如 Cutout）进行改进，因此误差高于 SAM 论文 中的结果。理论上，如果训练时间更长（例如 1800 轮而不是 200 轮），可以获得更低的误差，因为 SAM 不太容易过拟合。SAM 使用 rho=0.05，而 ASAM 则按照其作者的建议设置为 rho=2.0，详见 GitHub 问题 #37。

优化器	测试错误率
SGD + 动量	3.20 %
SAM + SGD + 动量	2.86 %
ASAM + SGD + 动量	2.55 %

引用

如果您在工作中使用此优化器，请引用原始作者：

@inproceedings{foret2021sharpnessaware,
  title={Sharpness-aware Minimization for Efficiently Improving Generalization},
  author={Pierre Foret and Ariel Kleiner and Hossein Mobahi and Behnam Neyshabur},
  booktitle={International Conference on Learning Representations},
  year={2021},
  url={https://openreview.net/forum?id=6Tm1mposlrM}
}

@inproceesings{pmlr-v139-kwon21b,
  title={ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks},
  author={Kwon, Jungmin and Kim, Jeongseop and Park, Hyunseo and Choi, In Kwon},
  booktitle ={Proceedings of the 38th International Conference on Machine Learning},
  pages={5905--5914},
  year={2021},
  editor={Meila, Marina and Zhang, Tong},
  volume={139},
  series={Proceedings of Machine Learning Research},
  month={18--24 Jul},
  publisher ={PMLR},
  pdf={http://proceedings.mlr.press/v139/kwon21b/kwon21b.pdf},
  url={https://proceedings.mlr.press/v139/kwon21b.html},
  abstract={近年来，以损失曲面的尖锐度作为泛化差距的有效度量的学习算法取得了最先进的性能。然而，在固定半径的刚性区域内定义的尖锐度，存在对参数缩放敏感的缺点，这会导致损失不变，从而削弱了尖锐度与泛化差距之间的联系。本文引入了尺度不变的自适应尖锐度概念，并提出了相应的泛化界。我们提出了一种新的学习方法——自适应尖锐度感知最小化（ASAM），并利用所提出的泛化界进行优化。在多个基准数据集上的实验结果表明，ASAM 显著提升了模型的泛化性能。}
}

SAM (Sharpness-Aware Minimization) 快速上手指南

SAM 是一种通过同时最小化损失值和损失锐度（loss sharpness）来提升模型泛化能力的优化器。它寻找具有统一低损失的参数邻域，能有效提升模型在各类数据集上的表现，并对标签噪声具有鲁棒性。本指南基于非官方 PyTorch 实现。

环境准备

• 系统要求: Linux, macOS, Windows
• Python 版本: Python 3.6+
• 核心依赖:
  • PyTorch >= 1.6.0
  • torchvision (可选，用于示例中的数据集加载)
• 前置知识: 熟悉 PyTorch 的基本训练循环流程。

安装步骤

该工具未发布到 PyPI，需直接从 GitHub 源码安装。

方式一：使用 pip 直接安装（推荐）

pip install git+https://github.com/davda54/sam.git

方式二：克隆仓库后安装（适合国内网络不稳定时手动下载）

如果直接访问 GitHub 较慢，可先手动下载代码包或使用国内镜像加速的 git 代理，然后执行：

git clone https://github.com/davda54/sam.git
cd sam
pip install -e .

基本使用

SAM 的核心机制是需要进行两次前向传播和反向传播（two forward-backward passes），因此训练速度约为普通优化器的一半。

1. 标准用法（分步执行）

这是最直观的使用方式，明确区分第一步（寻找最差点）和第二步（更新权重）。

from sam import SAM
import torch

# 初始化模型
model = YourModel()

# 定义基础优化器 (如 SGD)
base_optimizer = torch.optim.SGD

# 初始化 SAM 优化器
# rho: 邻域大小，默认 0.05；adaptive: 是否启用自适应版本 (ASAM)
optimizer = SAM(model.parameters(), base_optimizer, lr=0.1, momentum=0.9)

for input, output in data:
    # --- 第一次前向 - 反向传播 ---
    # 计算损失并反向传播
    loss = loss_function(output, model(input))
    loss.backward()
    # 执行第一步：找到局部邻域内损失最大的权重方向
    optimizer.first_step(zero_grad=True)
    
    # --- 第二次前向 - 反向传播 ---
    # 再次计算损失并反向传播 (确保完成完整的前向过程)
    loss_function(output, model(input)).backward()
    # 执行第二步：根据第一步找到的方向更新原始权重
    optimizer.second_step(zero_grad=True)

2. 进阶用法（闭包模式）

如果你习惯使用类似 LBFGS 的闭包风格，可以使用单步调用模式：

from sam import SAM

model = YourModel()
base_optimizer = torch.optim.SGD
optimizer = SAM(model.parameters(), base_optimizer, lr=0.1, momentum=0.9)

for input, output in data:
    def closure():
        loss = loss_function(output, model(input))
        loss.backward()
        return loss

    # 单次调用完成两步操作
    optimizer.step(closure)
    optimizer.zero_grad()

关键注意事项

• 批归一化 (Batch Normalization): 由于进行了两次前向传播，BN 层的运行统计量会被计算两次，这可能导致问题。建议在第二次前向传播时禁用运行统计量的更新。

  # 第一次传播：启用统计量更新
  enable_running_stats(model) 
  loss = loss_function(output, model(input))
  loss.backward()
  optimizer.first_step(zero_grad=True)
  
  # 第二次传播：禁用统计量更新
  disable_running_stats(model)
  loss_function(output, model(input)).backward()
  optimizer.second_step(zero_grad=True)
  

• 多 GPU 训练: 如果在多卡环境下训练，需在第一次反向传播时使用 no_sync() 上下文管理器，以避免梯度在不同卡之间过早同步。

  with model.no_sync():
      loss.backward()
  optimizer.first_step(zero_grad=True)
  

• 学习率调度: 学习率调度器应绑定到 base_optimizer 上，而不是 SAM 包装器本身。

  scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer.base_optimizer, T_max=200)
  

• ASAM (自适应版本): 若启用 adaptive=True，建议将 rho 参数调大至约 2.0（原文建议比标准 SAM 大 10 倍左右），通常能获得更好的效果。