如何使用 TorchVision 最新原语训练最先进的模型

由瓦西利斯·弗里尼奥蒂斯

几周前，TorchVision v0.11 版本发布，其中包含了许多新的原语、模型和训练食谱的改进，这些改进使得达到最先进（SOTA）的结果成为可能。该项目被称为“内置电池的 TorchVision”，旨在使我们的库现代化。我们希望研究人员能够通过使用常见的构建块更容易地重现论文和进行研究。此外，我们渴望为应用机器学习从业者提供必要的工具，让他们能够使用与研究中相同的最先进技术在自己的数据上训练模型。最后，我们希望刷新我们的预训练权重，并为用户提供更好的现成模型，希望他们能够构建更好的应用程序。

尽管还有很多工作要做，但我们想与您分享一些上述工作的令人兴奋的结果。我们将展示如何使用 TorchVision 中包含的新工具在 ResNet50 [1]这样高度竞争和深入研究架构上实现最先进的结果。我们将分享提高我们的基线超过 4.7 个准确度点以达到最终 top-1 准确率 80.9%的确切配方，并分享推导新训练过程的过程。此外，我们还将展示这个配方很好地推广到其他模型变体和家族。我们希望上述内容将影响未来的研究，以开发更强的可推广训练方法，并激励社区采用并贡献我们的努力。

结果

使用我们在 ResNet50 上找到的新训练配方，我们刷新了以下模型的预训练权重：

注意，除了 RetNet50 之外的所有模型的准确度都可以通过稍微调整它们的训练参数来进一步提高，但我们的重点是拥有一个适用于所有模型的单一稳健配方。

更新：我们已经刷新了 TorchVision 中大多数流行的分类模型，您可以在这篇博客文章中找到详细信息。

目前有两种方式可以使用该模型的最新权重。

使用多预训练权重 API

我们目前正在开发一个新的原型机制，该机制将扩展 TorchVision 的模型构建方法以支持多个权重。除了权重外，我们还存储有用的元数据（如标签、准确率、食谱链接等）以及使用模型所需的预处理转换。示例：

  from PIL import Image
  from torchvision import prototype as P
  img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")
   
  # Initialize model
  weights = P.models.ResNet50_Weights.IMAGENET1K_V2
  model = P.models.resnet50(weights=weights)
  model.eval()

  # Initialize inference transforms
  preprocess = weights.transforms()
   
  # Apply inference preprocessing transforms
  batch = preprocess(img).unsqueeze(0)
  prediction = model(batch).squeeze(0).softmax(0)
   
  # Make predictions
  label = prediction.argmax().item()
  score = prediction[label].item()
   
  # Use meta to get the labels
  category_name = weights.meta['categories'][label]
  print(f"{category_name}: {100 * score}%")

使用旧版 API

不想使用原型 API 的用户可以选择通过以下方法使用新权重：通过旧版 API 访问

  from torchvision.models import resnet
   
  # Overwrite the URL of the previous weights
  resnet.model_urls["resnet50"] = "https://download.pytorch.org/models/resnet50-11ad3fa6.pth"
   
  # Initialize the model using the legacy API
  model = resnet.resnet50(pretrained=True)
   
  # TODO: Apply preprocessing + call the model
  # ...

培训配方

我们的目的是利用 TorchVision 新引入的原始操作，推导出一个新的强大训练配方，该配方在从零开始训练 ImageNet 上的 vanilla ResNet50 架构时，实现了最先进的性能，且没有使用任何外部数据。尽管通过使用针对特定架构的技巧[2]可以进一步提高准确率，但我们决定不包含它们，以便该配方可以在其他架构中使用。我们的配方高度注重简洁性，并建立在 FAIR[3]、[4]、[5]、[6]、[7]的工作之上。我们的发现与 Wightman 等人[7]的并行研究相一致，他们也报告了通过关注训练配方实现了显著的准确率提升。

不再拖延，以下是我们的配方的主要参数：

  # Optimizer & LR scheme
  ngpus=8,
  batch_size=128,  # per GPU

  epochs=600, 
  opt='sgd',  
  momentum=0.9,

  lr=0.5, 
  lr_scheduler='cosineannealinglr', 
  lr_warmup_epochs=5, 
  lr_warmup_method='linear', 
  lr_warmup_decay=0.01, 


  # Regularization and Augmentation
  weight_decay=2e-05, 
  norm_weight_decay=0.0,

  label_smoothing=0.1, 
  mixup_alpha=0.2, 
  cutmix_alpha=1.0, 
  auto_augment='ta_wide', 
  random_erase=0.1, 
  
  ra_sampler=True,
  ra_reps=4,


  # EMA configuration
  model_ema=True, 
  model_ema_steps=32, 
  model_ema_decay=0.99998, 


  # Resizing
  interpolation='bilinear', 
  val_resize_size=232, 
  val_crop_size=224, 
  train_crop_size=176,

使用我们的标准训练参考脚本，我们可以使用以下命令训练 ResNet50：

torchrun --nproc_per_node=8 train.py --model resnet50 --batch-size 128 --lr 0.5 \
--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.00002 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--train-crop-size 176 --model-ema --val-resize-size 232 --ra-sampler --ra-reps 4

方法论

在我们的探索过程中，我们始终牢记以下原则：

1. 训练是一个随机过程，我们试图优化的验证指标是一个随机变量。这是由于采用的随机权重初始化方案以及在训练过程中存在的随机效应。这意味着我们不能通过单次运行来评估食谱更改的影响。标准做法是进行多次运行（通常 3 到 5 次）并研究总结统计（如平均值、标准差、中位数、最大值等）。
2. 通常不同参数之间存在显著的交互作用，特别是对于专注于正则化和减少过拟合的技术。因此，改变一个参数的值可能会影响其他参数的最佳配置。为了解决这个问题，可以采用贪婪搜索方法（这通常会导致次优结果但实验可行）或应用网格搜索（这会导致更好的结果但计算成本高昂）。在本工作中，我们使用了这两种方法的混合。
3. 非确定性的技术或引入噪声的技术通常需要更长的训练周期来提高模型性能。为了保持可追踪性，我们最初使用了短训练周期（少量 epoch）来决定哪些路径可以提前消除，哪些应该使用更长的训练来探索。
4. 由于重复实验，存在过度拟合验证数据集的风险[8]。为了减轻一些风险，我们只应用提供显著精度改进的训练优化，并使用 K 折交叉验证来验证对验证集进行的优化。此外，我们还确认我们的配方成分对其他我们没有优化超参数的模型也具有良好的泛化能力。

关键精度改进的分解

如前文博客文章所述，训练模型并非单调递增准确率的旅程，过程中涉及大量回溯。为了量化每次优化的效果，以下我们尝试展示从 TorchVision 原始配方开始的理想化线性旅程，以推导出最终配方。我们想明确指出，这实际上是实际路径的过度简化，因此应持保留态度。

下表提供了在基线之上叠加增量改进的性能总结。除非另有说明，否则我们报告的是 3 次运行中 Acc@1 最佳的模型：

*在最后一个模型的基础上进行了推理尺寸的调整。详情见下文。

** 文章发布后由社区贡献的内容。详情见下文。

基准

我们的基线是之前发布的 TorchVision 的 ResNet50 模型。其训练过程如下：

  # Optimizer & LR scheme
  ngpus=8,
  batch_size=32,  # per GPU

  epochs=90, 
  opt='sgd',  
  momentum=0.9,

  lr=0.1, 
  lr_scheduler='steplr', 
  lr_step_size=30, 
  lr_gamma=0.1, 


  # Regularization
  weight_decay=1e-4,


  # Resizing
  interpolation='bilinear', 
  val_resize_size=256, 
  val_crop_size=224, 
  train_crop_size=224,

上述大部分参数都是我们训练脚本中的默认值。我们将在此基础上进行优化，逐步达到最终的配方。

学习率优化

我们可以应用一些参数更新来提高训练的准确性和速度。这可以通过增加批量大小和调整学习率（LR）来实现。另一种常见的方法是应用预热并逐渐增加我们的学习率。这特别有利于我们使用非常高的学习率时，有助于早期 epoch 的训练稳定性。最后，另一种优化是应用余弦调度来调整 epoch 期间的学习率。余弦的一个大优点是无需优化超参数，这减少了我们的搜索空间。

以下是基于基线方案的额外优化。请注意，我们已经进行了多次实验以确定参数的最佳配置：

  batch_size=128,  # per GPU

  lr=0.5, 
  lr_scheduler='cosineannealinglr', 
  lr_warmup_epochs=5, 
  lr_warmup_method='linear', 
  lr_warmup_decay=0.01,

与基线相比，上述优化将我们的 top-1 准确率提高了 0.364 个百分点。请注意，为了结合不同的 LR 策略，我们使用了新引入的 SequentialLR 调度器。

简单增强

原始模型使用基本的增强变换进行训练，例如随机裁剪和水平翻转。提高我们准确度的一个简单方法是应用更复杂的“自动增强”技术。对我们来说表现最好的是 TrivialAugment [9]，它极其简单，可以被认为是“参数自由”，这意味着它可以进一步帮助我们缩小搜索空间。

这里是应用在上一个步骤之上的更新：

auto_augment='ta_wide',

使用 TrivialAugment 使我们的 top-1 准确率相比之前提高了 0.312 个百分点。

长时间训练

长的训练周期在包含行为随机的成分的配方中是有益的。更具体地说，当我们开始添加越来越多引入噪声的技术时，增加训练轮数变得至关重要。请注意，在我们探索的早期阶段，我们使用了大约 200 个轮次的相对较短的周期，后来增加到 400 个轮次，当我们开始缩小大多数参数时，最终在配方的最终版本中增加到 600 个轮次。

下面我们看到在早期步骤之上应用了更新：

epochs=600,

这进一步提高了我们基于前一步骤的 top-1 准确率 1.8 个百分点。这是我们在此迭代过程中观察到的最大增幅。值得注意的是，这种单一优化的效果被夸大了，并且有些误导。仅仅增加旧基线上的训练轮数并不会带来如此显著的改进。尽管如此，LR 优化与强大的增强策略的结合有助于模型从更长的周期中受益。还值得一提的是，我们在早期过程中引入长时间训练周期的原因是因为在接下来的步骤中，我们将引入需要更多轮数才能提供良好结果的技巧。

随机擦除

另一种已知有助于提高分类准确率的数据增强技术是随机擦除[10]，[11]。通常与自动增强方法结合使用，它通常由于其正则化效果而带来额外的准确率提升。在我们的实验中，我们仅通过网格搜索调整了应用该方法的概率，并发现将其概率保持在低水平，通常在 10%左右，是有益的。

在此之上引入了额外的参数：

random_erase=0.1,

应用随机擦除可以将我们的 Acc@1 提高 0.190 个点。

标签平滑

为了减少过拟合，一个好的技术是阻止模型变得过于自信。这可以通过使用标签平滑[12]来软化真实标签来实现。有一个参数控制平滑的程度（值越高越强），我们需要指定。虽然可以通过网格搜索来优化它，但我们发现 0.05-0.15 的值会产生相似的结果，因此为了避免过拟合，我们使用了它在引入该技术的论文中使用的相同值。

在此步骤中我们可以找到添加的额外配置：

label_smoothing=0.1,

我们使用了 PyTorch 新引入的 CrossEntropyLoss 标签平滑参数，这使我们的准确率额外提高了 0.318 个百分点。

Mixup 和 Cutmix

两种常用于产生 SOTA 结果的数据增强技术是 Mixup 和 Cutmix [13]，[14]。它们都通过软化标签和图像来提供强大的正则化效果。在我们的设置中，我们发现随机应用其中之一具有同等概率是有益的。每个都使用超参数 alpha 进行参数化，该参数控制从其中采样平滑概率的 Beta 分布的形状。我们进行了非常有限的网络搜索，主要关注论文中提出的常见值。

以下为您提供了两种技术的 alpha 参数的最佳值：

mixup_alpha=0.2, 
cutmix_alpha=1.0,

应用 mixup 可以使我们的准确率提高 0.118 个点，与 cutmix 结合使用则额外提高 0.278 个点。

权重衰减调整

我们的标准配方使用 L2 正则化来减少过拟合。权重衰减参数控制正则化的程度（越大越强），默认情况下应用于模型的所有学习参数。在此配方中，我们对标准方法进行了两项优化。首先，我们执行网格搜索来调整权重衰减参数；其次，我们禁用了归一化层参数的权重衰减。

以下可以找到我们食谱中权重衰减的最佳配置：

weight_decay=2e-05, 
norm_weight_decay=0.0,

上述更新使我们的准确率提高了 0.526 个点，为已知事实提供了额外的实验证据，即调整权重衰减对模型性能有显著影响。我们将归一化参数与其他参数分离的方法受到了 ClassyVision 方法的启发。

FixRes 缓解措施

在我们的实验初期，我们发现了一个重要特性，即如果验证时使用的分辨率从训练的 224x224 提高到更高，模型的表现会显著更好。这一效应在 FixRes 论文[5]中进行了详细研究，并提出了两种缓解措施：a) 可以尝试降低训练分辨率，以使验证分辨率上的准确率最大化；或者 b) 可以在两阶段训练中微调模型，使其适应目标分辨率。由于我们不希望引入两阶段训练，所以我们选择了方案 a）。这意味着我们将训练裁剪大小从 224 降低，并使用网格搜索找到最大化 224x224 分辨率验证的方案。

以下是我们配方中使用的最佳值：

val_crop_size=224, 
train_crop_size=176,

上述优化使我们的准确率提高了 0.160 个点，并将训练速度提高了 10%。

值得注意的是，FixRes 效果仍然存在，这意味着当我们提高分辨率时，模型在验证集上的表现仍然更好。此外，进一步减少训练裁剪尺寸实际上会损害准确率。从直观上看，这是因为当分辨率降低到一定程度后，图像中的关键细节开始消失。最后，我们应该指出，上述 FixRes 缓解似乎对与 ResNet50 深度相似度的模型有益。具有更大感受野的更深变体似乎略有负面影响（通常为 0.1-0.2 个点）。因此，我们将这部分配方视为可选。以下我们可视化了最佳可用检查点（带有完整配方）的性能，这些检查点是使用 176 和 224 分辨率训练的模型：

指数移动平均（EMA）

EMA 是一种技术，允许在不增加模型复杂度或推理时间的情况下提高模型的准确性。它对模型权重执行指数移动平均，这导致准确性提高和模型更稳定。平均每几轮迭代发生一次，其衰减参数通过网格搜索进行调整。

下面可以看到我们食谱的最佳值：

model_ema=True, 
model_ema_steps=32, 
model_ema_decay=0.99998,

使用 EMA 将我们的准确性提高了 0.254 个点，与上一步相比。请注意，TorchVision 的 EMA 实现建立在 PyTorch 的 AveragedModel 类之上，关键区别在于它不仅平均模型参数，还平均其缓冲区。此外，我们还采用了 Pycls 的技巧，以使衰减参数的参数化不依赖于 epoch 的数量。

推理调整大小

与其他所有涉及使用不同参数训练模型的步骤不同，这次优化是在最终模型之上进行的。在推理过程中，图像被调整到特定分辨率，然后从中裁剪出一个中心 224x224 的区域。原始配方使用了 256 的调整大小，这导致了与 FixRes 论文[5]中描述的类似差异。通过将此调整大小值更接近目标推理分辨率，可以提高准确性。为了选择值，我们在[224, 256]区间内以 8 为步长进行短暂的网格搜索。为了避免过拟合，使用验证集的一半来选择值，并使用另一半来确认。

下面是我们在配方中使用的最优值：

val_resize_size=232,

上述优化提高了我们的准确度 0.224 个百分点。值得注意的是，ResNet50 的最优值对于 ResNet101、ResNet152 和 ResNeXt50 也效果最佳，这表明它可以在不同模型之间泛化：

[更新] 重复增强

重复增强[15]，[16]是另一种可以提高整体准确率的技巧，已被其他强大配方如[6]，[7]所采用。社区贡献者 Tal Ben-Nun 通过提出使用 4 次重复训练模型来进一步改进我们的原始配方。他的贡献是在这篇文章发布之后。

下面您可以看到我们配方中使用的最优值：

ra_sampler=True,
ra_reps=4,

上述是最终的优化，提高了我们的准确率 0.184 个百分点。

测试过但未采用的优化

在我们研究的早期阶段，我们尝试了额外的技术、配置和优化。由于我们的目标是尽可能保持我们的配方简单，我们决定不包括任何没有带来显著改进的内容。以下是我们采取的一些方法，但它们没有进入我们的最终配方：

• 优化器：使用更复杂的优化器，如 Adam、RMSProp 或具有 Nesterov 动量的 SGD，与具有动量的普通 SGD 相比，并没有带来显著更好的结果。
• 学习率调度器：我们尝试了不同的学习率调度器方案，如 StepLR 和指数调度器。尽管后者与 EMA 一起使用时效果更好，但它通常需要额外的超参数，例如定义最小学习率才能正常工作。相反，我们只使用余弦退火将学习率衰减到零，并选择具有最高精度的检查点。
• 自动增强：我们尝试了不同的增强策略，如 AutoAugment 和 RandAugment。这些方法都没有优于更简单的参数免费 TrivialAugment。
• 插值：使用双三次插值或最近邻插值并没有比双线性插值提供显著更好的结果。
• 归一化层：使用同步批量归一化并没有比使用常规批量归一化带来显著更好的结果。

致谢

我们要感谢 Piotr Dollar、Mannat Singh 和 Hugo Touvron 在开发食谱期间提供他们的见解和反馈，以及他们之前关于我们食谱基础的研究工作。他们的支持对于实现上述结果是无价的。此外，我们还要感谢 Prabhat Roy、Kai Zhang、Yiwen Song、Joel Schlosser、Ilqar Ramazanli、Francisco Massa、Mannat Singh、Xiaoliang Dai、Samuel Gabriel、Allen Goodman 和 Tal Ben-Nun 对“电池包含”项目的贡献。

参考文献列表

1. 何凯明，张祥宇，任少卿，孙剑。“深度残差学习在图像识别中的应用”。
2. 唐和，张志，张航，张中月，谢俊元，李穆. “卷积神经网络在图像分类中的技巧集”
3. 皮奥特·多拉尔，曼纳特·辛格，罗斯·吉里克. “快速且精确的模型缩放”
4. 肖铁，曼纳特·辛格，埃里克·明顿，特雷弗·达雷尔，皮奥特·多拉尔，罗斯·吉里克. “早期卷积有助于 Transformer 更好地观察”
5. 于戈·图尔沃，安德烈亚·韦达利，马蒂吉斯·杜兹，赫维·热戈. “修复训练-测试分辨率差异问题”
6. 雨果·图尔沃恩，马蒂厄·科尔，马蒂吉斯·杜兹，弗朗西斯科·马斯，亚历山大·萨布罗莱斯，赫维·热戈。 “训练数据高效图像变换器与注意力蒸馏”
7. 罗斯·怀特曼，雨果·图尔沃恩，赫维·热戈。 “ResNet 反击：timm 中的改进训练过程”
8. 本杰明·雷希特，丽贝卡·罗尔夫斯，路德维希·施密特，瓦伊沙尔·桑卡尔。 “ImageNet 分类器是否泛化到 ImageNet？”
9. 萨缪尔·G·穆勒，弗兰克·胡特。 “TrivialAugment：无需调整且达到最先进水平的数据增强”
10. 钟骏，梁正，康国梁，李少子，杨毅. “随机擦除数据增强”
11. 特伦斯·德弗里斯，格雷厄姆·W·泰勒. “使用 Cutout 改进卷积神经网络的正则化”
12. 克里斯蒂安·塞格迪，文森特·范豪克，谢尔盖·约费，乔恩·舒伦斯，兹比格涅夫·沃伊纳. “重新思考计算机视觉中的 Inception 架构”
13. 张宏毅，穆斯塔法·西塞，扬·N·道宾，大卫·洛佩斯-帕斯. “mixup：超越经验风险最小化”
14. Sangdoo Yun，Dongyoon Han，Seong Joon Oh，Sanghyuk Chun，Junsuk Choe，Youngjoon Yoo. “CutMix：具有可定位特征的强大分类器的正则化策略”
15. Elad Hoffer，Tal Ben-Nun，Itay Hubara，Niv Giladi，Torsten Hoefler，Daniel Soudry. “扩大你的批量：使用更大的批量进行更好的训练”
16. Maxim Berman，Hervé Jégou，Andrea Vedaldi，Iasonas Kokkinos，Matthijs Douze. “Multigrain：统一图像嵌入用于类别和实例”