跳转到主要内容
博客

如何使用 TorchVision 最新基元训练最先进的模型

作者: 2021 年 11 月 18 日2024 年 11 月 15 日暂无评论

几周前,TorchVision v0.11 发布,其中包含众多新原语、模型和训练秘籍改进,从而能够实现最先进 (SOTA) 的结果。该项目被命名为“内置电池的 TorchVision”,旨在使我们的库现代化。我们希望通过使用常见的构建块,使研究人员能够更轻松地复现论文和进行研究。此外,我们渴望为应用机器学习从业者提供必要的工具,使其能够使用与研究中相同的 SOTA 技术在自己的数据上训练模型。最后,我们希望更新我们的预训练权重,并为用户提供更好的开箱即用模型,希望他们能构建更好的应用程序。

尽管仍有许多工作要做,但我们想与您分享上述工作的一些令人兴奋的结果。我们将展示如何使用 TorchVision 中包含的新工具,在 ResNet50 [1] 这种竞争激烈且经过充分研究的架构上实现最先进的结果。我们将分享用于将我们的基线提高 4.7 个准确率点,以达到最终 80.9% 的 top-1 准确率的精确秘籍,并分享推导新训练过程的历程。此外,我们将展示该秘籍对其他模型变体和系列也具有良好的泛化能力。我们希望上述内容将影响未来开发更强大的可泛化训练方法的研究,并激励社区采纳和贡献我们的工作。

结果

使用我们新发现的 ResNet50 训练秘籍,我们已更新以下模型的预训练权重

模型 准确率@1 准确率@5
ResNet50 80.858 95.434
ResNet101 81.886 95.780
ResNet152 82.284 96.002
ResNeXt50-32x4d 81.198 95.340

请注意,除了 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] 可以进一步提高准确率,但我们决定不包含它们,以便该秘籍可以在其他架构中使用。我们的秘籍[3],[4],[5],[6],[7] 大力关注简单性,并基于 FAIR 的工作构建。我们的发现与 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 原始秘籍开始推导最终秘籍的理想化线性过程。我们想澄清的是,这过度简化了我们实际遵循的路径,因此应持保留态度。

Cumulative Accuracy Improvements for ResNet50

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

准确率@1 准确率@5 增量差异 绝对差异
ResNet50 基线 76.130 92.862 0.000 0.000
+ 学习率优化 76.494 93.198 0.364 0.364
+ TrivialAugment 76.806 93.272 0.312 0.676
+ 长期训练 78.606 94.052 1.800 2.476
+ 随机擦除 78.796 94.094 0.190 2.666
+ 标签平滑 79.114 94.374 0.318 2.984
+ Mixup 79.232 94.536 0.118 3.102
+ Cutmix 79.510 94.642 0.278 3.380
+ 权重衰减调整 80.036 94.746 0.526 3.906
+ FixRes 缓解 80.196 94.672 0.160 4.066
+ EMA 80.450 94.908 0.254 4.320
+ 推理尺寸调整 * 80.674 95.166 0.224 4.544
+ 重复增强 ** 80.858 95.434 0.184 4.728

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

** 文章发布后进行的社区贡献。详情请见下文。

基线

我们的基线是之前发布的 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,

上述大部分参数是我们训练脚本中的默认值。我们将在此基线之上逐步引入优化,直到最终得到最终秘籍。

学习率优化

我们可以应用一些参数更新来提高训练的准确性和速度。这可以通过增加批量大小和调整学习率来实现。另一种常见方法是应用预热并逐步增加学习率。这在学习率非常高时特别有益,有助于训练在早期 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 个点。请注意,为了结合不同的学习率策略,我们使用了新引入的 SequentialLR 调度器。

TrivialAugment

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

以下是在上一步基础上应用的更新:

auto_augment='ta_wide',

与上一步相比,使用 TrivialAugment 将我们的 top-1 准确率提高了 0.312 个百分点。

长期训练

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

下面我们看到在前面步骤的基础上应用的更新:

epochs=600,

这在之前步骤的基础上进一步将我们的 top-1 准确率提高了 1.8 个点。这是我们在此迭代过程中将观察到的最大增幅。值得注意的是,这种单一优化的效果被夸大了,并有些误导。仅仅在旧基线之上增加 epoch 数量并不会产生如此显著的改进。然而,学习率优化与强大增强策略的结合有助于模型从更长的周期中获益。还值得一提的是,我们如此早在过程中引入冗长的训练周期的原因是,在接下来的步骤中,我们将引入需要显著更多 epoch 才能提供良好结果的技术。

随机擦除

另一种有助于分类准确性的数据增强技术是随机擦除 [10][11]。它通常与自动增强方法配合使用,由于其正则化效果,通常会进一步提高准确性。在我们的实验中,我们仅通过网格搜索调整了应用该方法的概率,发现将其概率保持在较低水平(通常约为 10%)是有益的。

以下是在之前基础上引入的额外参数:

random_erase=0.1,

应用随机擦除将我们的 Acc@1 进一步提高了 0.190 个点。

标签平滑

减少过拟合的一个好方法是阻止模型过于自信。这可以通过使用标签平滑来软化真实标签来实现[12]。有一个单一参数控制平滑的程度(越高越强),我们需要指定它。尽管通过网格搜索优化它是可能的,但我们发现大约 0.05-0.15 的值会产生类似的结果,因此为了避免过拟合,我们使用了与引入它的论文中相同的值。

下面我们可以找到此步骤中添加的额外配置:

label_smoothing=0.1,

我们使用 PyTorch 新引入的 CrossEntropyLoss 的 label_smoothing 参数,这使我们的准确率额外提高了 0.318 个点。

Mixup 和 Cutmix

Mixup 和 Cutmix [13],[14] 是两种常用于生成 SOTA 结果的数据增强技术。它们都通过不仅软化标签而且软化图像来提供强大的正则化效果。在我们的设置中,我们发现以相同的概率随机应用其中一种是有益的。每种技术都由一个超参数 alpha 进行参数化,该参数控制贝塔分布的形状,从该分布中抽取平滑概率。我们进行了非常有限的网格搜索,主要关注论文中提出的常见值。

您将在下面找到两种技术的 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 缓解措施

在我们的实验早期发现的一个重要特性是,如果验证期间使用的分辨率从训练的 224×224 增加,模型的性能会显著提高。FixRes 论文 [5] 详细研究了这种效应,并提出了两种缓解措施:a) 可以尝试降低训练分辨率,以使验证分辨率上的准确率最大化;或者 b) 可以通过两阶段训练对模型进行微调,使其适应目标分辨率。由于我们不想引入两阶段训练,因此我们选择了选项 a)。这意味着我们降低了训练裁剪尺寸,从 224 开始,并使用网格搜索找到在 224×224 分辨率下验证准确率最大化的尺寸。

您可以在下面看到我们的秘籍中使用的最佳值:

val_crop_size=224, 
train_crop_size=176,

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

值得注意的是,FixRes 效应仍然存在,这意味着当我们在验证时增加分辨率时,模型性能仍然更好。此外,进一步减小训练裁剪尺寸实际上会损害准确率。这从直觉上是说得通的,因为在关键细节开始从图片中消失之前,分辨率只能降低这么多。最后,我们应该指出,上述 FixRes 缓解措施似乎对深度与 ResNet50 相似的模型有益。具有较大感受野的更深变体似乎受到轻微负面影响(通常为 0.1-0.2 个点)。因此,我们认为这部分秘籍是可选的。下面我们可视化了使用 176 和 224 分辨率训练的模型(使用完整秘籍)的最佳可用检查点的性能。

Best ResNet50 trained with 176 Resolution
Best ResNet50 trained with 224 Resolution

指数移动平均 (EMA)

EMA 是一种可以提高模型准确率而又不增加其复杂性或推理时间的技术。它对模型权重执行指数移动平均,从而提高准确率并使模型更稳定。平均每隔几次迭代进行一次,其衰减参数通过网格搜索进行调整。

您可以在下面看到我们秘籍的最佳值:

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

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

推理尺寸调整

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

您可以在下面看到我们的秘籍中使用的最佳值:

val_resize_size=232,

上述优化使我们的准确率提高了 0.224 个百分点。值得注意的是,ResNet50 的最佳值也最适用于 ResNet101、ResNet152 和 ResNeXt50,这暗示它在模型之间具有泛化性。

ResNet50 Inference Resize
ResNet101 Inference Resize
Best ResNet50 trained with 224 Resolution

[更新] 重复增强

重复增强 [15],[16] 是另一种可以提高整体准确率的技术,并已被其他强大的秘籍使用,例如 [6],[7]。社区贡献者 Tal Ben-Nun 进一步改进 了我们的原始秘籍,提议以 4 次重复训练模型。他的贡献是在本文发布之后。

您可以在下面看到我们的秘籍中使用的最佳值:

ra_sampler=True,
ra_reps=4,

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

已测试但未采用的优化

在研究的早期阶段,我们尝试了额外的技术、配置和优化。由于我们的目标是尽可能简化我们的秘籍,我们决定不包含任何未提供显著改进的内容。以下是我们尝试过但未纳入最终秘籍的一些方法:

  • 优化器: 使用更复杂的优化器,如 Adam、RMSProp 或带有 Nesterov 动量的 SGD,并未比普通的带有动量的 SGD 提供显著更好的结果。
  • 学习率调度器: 我们尝试了不同的学习率调度器方案,如 StepLR 和 Exponential。尽管后者与 EMA 配合效果更好,但它通常需要额外的超参数,例如定义最低学习率才能良好工作。相反,我们只是使用余弦退火将学习率衰减到零,并选择准确率最高的检查点。
  • 自动增强: 我们尝试了不同的增强策略,如 AutoAugment 和 RandAugment。这些都没有超越更简单的无参数 TrivialAugment。
  • 插值: 使用双三次或最近邻插值并未比双线性插值提供显著更好的结果。
  • 归一化层: 使用 Sync Batch Norm 未比使用常规 Batch Norm 产生显著更好的结果。

致谢

我们衷心感谢 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. Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, Mu Li。“卷积神经网络图像分类的技巧包”
  3. Piotr Dollár, Mannat Singh, Ross Girshick. “快速准确的模型缩放”
  4. Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, Ross Girshick. “早期卷积帮助 Transformer 更好地观察”
  5. Hugo Touvron, Andrea Vedaldi, Matthijs Douze, Hervé Jégou. “修复训练-测试分辨率差异”
  6. Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou. “通过注意力训练数据高效的图像 Transformer & 蒸馏”
  7. Ross Wightman, Hugo Touvron, Hervé Jégou. “ResNet 卷土重来:timm 中改进的训练过程”
  8. Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, Vaishaal Shankar. “ImageNet 分类器能否泛化到 ImageNet?”
  9. Samuel G. Müller, Frank Hutter. “TrivialAugment:无需调参却最先进的数据增强”
  10. Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, Yi Yang. “随机擦除数据增强”
  11. Terrance DeVries, Graham W. Taylor. “通过 Cutout 改进卷积神经网络的正则化”
  12. Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, Zbigniew Wojna. “重新思考计算机视觉的 Inception 架构”
  13. Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz. “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. “多粒度:用于类别和实例的统一图像嵌入”