几周前，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 |
请注意，除了 ResNet50 之外，所有模型的准确性都可以通过稍微调整其训练参数来进一步提高，但我们的重点是拥有一套适用于所有模型的单一稳健方案。
更新：我们已经刷新了 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 上从头开始训练，不使用额外的外部数据时，能为香草 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
方法论
在我们的探索过程中，我们牢记以下几个原则
- 训练是一个随机过程，我们试图优化的验证指标是一个随机变量。这是由于所采用的随机权重初始化方案和训练过程中随机效应的存在。这意味着我们不能通过单次运行来评估方案更改的效果。标准做法是进行多次运行（通常为 3 到 5 次）并研究汇总统计数据（例如平均值、标准差、中位数、最大值等）。
- 不同参数之间通常存在显著的相互作用，特别是对于侧重于正则化和减少过拟合的技术。因此，更改其中一个参数的值可能会影响其他参数的最佳配置。为了解决这个问题，可以采用贪婪搜索方法（通常会导致次优结果但可追踪的实验）或应用网格搜索（会导致更好的结果但计算成本高）。在这项工作中，我们混合使用了这两种方法。
- 非确定性或引入噪声的技术通常需要更长的训练周期才能提高模型性能。为了使事情易于处理，我们最初使用短训练周期（少量 epoch）来决定哪些路径可以提前消除，哪些应该使用更长的训练周期进行探索。
- 由于重复实验，存在验证数据集过拟合的风险 [8]。为了减轻部分风险，我们只应用提供显著准确性改进的训练优化，并使用 K 折交叉验证来验证在验证集上进行的优化。此外，我们确认我们的方案成分在其他模型上也能很好地推广，而我们并未优化这些模型的超参数。
关键准确性改进细分
正如早期博客文章中所讨论的，训练模型并非一个准确性单调递增的旅程，该过程涉及大量的回溯。为了量化每次优化的效果，下面我们尝试展示一个从 TorchVision 原始方案开始推导最终方案的理想化线性旅程。我们想澄清的是，这过度简化了我们实际遵循的路径，因此应慎重对待。
在下表中，我们总结了在基线之上叠加增量改进的性能。除非另有说明，否则我们报告 3 次运行中 Acc@1 最佳的模型
| 准确率@1 | 准确率@5 | 增量差异 | 绝对差异 | |
|---|---|---|---|---|
| ResNet50 基线 | 76.130 | 92.862 | 0.000 | 0.000 |
| + LR 优化 | 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,
上述大多数参数都是我们训练脚本中的默认值。我们将在此基线之上逐步引入优化，直到最终达到最终方案。
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 个点。请注意，为了结合不同的学习率策略，我们使用了新引入的 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 标签平滑参数，这使我们的准确率额外提高了 0.318 个点。
Mixup 和 Cutmix
Mixup 和 Cutmix [13]、[14] 是两种常用于产生 SOTA 结果的数据增强技术。它们都通过不仅软化标签而且软化图像来提供强大的正则化效果。在我们的设置中，我们发现以相同的概率随机应用其中一种是有益的。每种技术都由一个超参数 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 缓解
在我们实验早期发现的一个重要特性是，如果验证期间使用的分辨率从训练的 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 分辨率训练的最佳可用检查点（包含完整方案）的性能
指数移动平均 (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，这表明它可以在模型之间推广
[更新] 重复增强
重复增强 [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。
- 插值：使用双三次或最近邻插值并未比双线性插值提供显著更好的结果。
- 归一化层：使用同步批量归一化并未比使用常规批量归一化产生显著更好的结果。
鸣谢
我们衷心感谢 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 对“内置电池”项目的贡献。
参考文献
- Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. “深度残差学习用于图像识别”。
- Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, Mu Li. “卷积神经网络图像分类的技巧包”
- Piotr Dollár, Mannat Singh, Ross Girshick. “快速准确的模型缩放”
- Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, Ross Girshick. “早期卷积帮助 Transformers 更好地观察”
- Hugo Touvron, Andrea Vedaldi, Matthijs Douze, Hervé Jégou. “修复训练-测试分辨率差异”
- Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou. “通过注意力训练数据高效图像 transformer 和蒸馏”
- Ross Wightman, Hugo Touvron, Hervé Jégou. “ResNet 卷土重来：timm 中改进的训练过程”
- Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, Vaishaal Shankar. “ImageNet 分类器能否推广到 ImageNet？”
- Samuel G. Müller, Frank Hutter. “TrivialAugment：免调优但最先进的数据增强”
- Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, Yi Yang. “随机擦除数据增强”
- Terrance DeVries, Graham W. Taylor. “通过 Cutout 改进卷积神经网络的正则化”
- Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, Zbigniew Wojna. “重新思考计算机视觉的 Inception 架构”
- Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz. “mixup：超越经验风险最小化”
- Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, Youngjoon Yoo. “CutMix：用于训练具有局部特征的强大分类器的正则化策略”
- Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, Daniel Soudry. “增强你的批次：用更大的批次进行更好的训练”
- Maxim Berman, Hervé Jégou, Andrea Vedaldi, Iasonas Kokkinos, Matthijs Douze. “Multigrain：用于类别和实例的统一图像嵌入”