几周前，TorchVision v0.11 正式发布，其中包含了大量新的原语、模型以及训练配方的改进，使得模型能够达到最先进（SOTA）的水平。该项目被戏称为“内置电池的 TorchVision”（TorchVision with Batteries Included），旨在实现我们库的现代化。我们希望通过使用通用的构建模块，让研究人员能够更轻松地复现论文并进行研究。此外，我们渴望提供必要的工具，帮助应用机器学习从业者使用与研究中相同的 SOTA 技术在自己的数据上训练模型。最后，我们希望更新预训练权重，为用户提供更好的现成模型，并期待他们能够构建出更出色的应用程序。
尽管还有很多工作要做，但我们还是想与大家分享上述工作中的一些令人兴奋的成果。我们将展示如何使用 TorchVision 中包含的新工具，在 ResNet50 [1] 这样竞争激烈且研究深入的架构上实现 SOTA 结果。我们将分享用于将基准性能提升超过 4.7 个准确率点、最终达到 80.9% Top-1 准确率的确切配方，并分享推导这一新训练过程的历程。此外，我们还将证明该配方可以很好地推广到其他模型变体和系列。我们希望上述内容能影响未来开发更强大、更具通用性训练方法的研究，并激励社区采纳我们的成果并为此做出贡献。
结果
使用我们在 ResNet50 上发现的新训练配方，我们更新了以下模型的预训练权重：
| 模型 | Accuracy@1 (Top-1 准确率) | Accuracy@5 (Top-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 上从零开始训练且不使用任何额外外部数据的情况下，为原生 ResNet50 架构实现 SOTA 结果。虽然通过使用架构特定的技巧 [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
方法论
在探索过程中，我们遵循了以下几项原则：
- 训练是一个随机过程，我们试图优化的验证指标是一个随机变量。这是由于采用了随机权重初始化方案以及训练过程中存在随机效应所致。这意味着我们不能仅通过单次运行来评估配方变更的影响。标准的做法是进行多次运行（通常为 3 到 5 次）并研究汇总统计数据（如平均值、标准差、中位数、最大值等）。
- 不同参数之间通常存在显著的相互作用，特别是对于侧重于正则化和减少过拟合的技术而言。因此，更改其中一个参数的值可能会对其他参数的最佳配置产生影响。为了解决这个问题，可以采用贪婪搜索方法（通常会导致次优结果但实验易于处理）或应用网格搜索（会导致更好的结果但计算昂贵）。在这项工作中，我们混合使用了这两种方法。
- 非确定性或引入噪声的技术通常需要更长的训练周期来提高模型性能。为了保持可处理性，我们最初使用较短的训练周期（较少的 epoch 数）来决定哪些路径可以尽早消除，哪些应该通过更长的训练进行探索。
- 由于重复实验，存在过拟合验证数据集 [8] 的风险。为了减轻这种风险，我们只应用能带来显著准确率提升的训练优化，并使用 K 折交叉验证来验证在验证集上完成的优化。此外，我们确认了我们的配方成分在其他未进行超参数优化的模型上也具有良好的泛化能力。
关键准确率改进的分解
正如在之前的博文中所讨论的那样，训练模型并不是一个准确率单调递增的过程，它涉及大量的回溯。为了量化每次优化的效果，我们在下文中尝试展示一个理想化的线性旅程，展示如何从原始的 TorchVision 配方推导出最终配方。我们想澄清的是，这只是我们实际所走路径的过度简化，因此请谨慎看待。
在下表中，我们提供了在基准基础上进行叠加增量改进后的性能摘要。除非另有说明，否则我们报告的是 3 次运行中 Acc@1 最好的模型。
| Accuracy@1 (Top-1 准确率) | Accuracy@5 (Top-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 |
| + 随机擦除 (Random Erasing) | 78.796 | 94.094 | 0.190 | 2.666 |
| + 标签平滑 (Label Smoothing) | 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 |
| + 权重衰减 (Weight Decay) 调整 | 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 |
| + 重复增强 (Repeated Augmentation) ** | 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 优化
我们可以应用一些参数更新来提高训练的准确率和速度。这可以通过增加批处理大小 (batch size) 和调整学习率 (LR) 来实现。另一种常见的方法是应用预热 (warmup) 并逐渐增加学习率。这在采用非常高的学习率时特别有益，有助于稳定早期 epoch 的训练。最后，另一个优化是应用余弦调度 (Cosine Schedule) 来在 epoch 期间调整 LR。余弦调度的一大优势是没有需要优化的超参数，这缩减了我们的搜索空间。
以下是在基准配方之上应用的额外优化。注意，我们运行了多次实验来确定参数的最佳配置：
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
原始模型是使用基础增强变换（如随机裁剪缩放和水平翻转）进行训练的。提高准确率的一种简单方法是应用更复杂的“自动增强”技术。对我们来说表现最好的是 TrivialAugment [9]，它非常简单，可以被视为“无参数”，这意味着它可以帮助我们进一步缩减搜索空间。
以下是在上一步基础上应用的更新：
auto_augment='ta_wide',
与上一步相比，使用 TrivialAugment 将我们的 Top-1 准确率提高了 0.312 个点。
长时间训练
当我们的配方包含随机行为的成分时，更长的训练周期是有益的。更具体地说，随着我们开始添加越来越多的引入噪声的技术，增加 epoch 的数量变得至关重要。注意，在我们探索的早期阶段，我们使用了大约 200 个 epoch 的相对较短的周期，后来随着我们开始缩小大多数参数的范围，增加到了 400 个，并在配方的最终版本中增加到了 600 个 epoch。
下面我们看到的是在之前步骤基础上应用的更新：
epochs=600,
这使我们的 Top-1 准确率在之前步骤的基础上进一步提高了 1.8 个点。这是我们在迭代过程中观察到的最大幅度提升。值得注意的是，这种单一优化的效果被夸大了，甚至有些误导。仅仅在旧基准上增加 epoch 数量并不会产生如此显著的提升。尽管如此，LR 优化与强力增强策略的结合有助于模型从更长的周期中受益。同样值得一提的是，我们之所以如此早地引入漫长的训练周期，是因为在接下来的步骤中，我们将引入需要更多 epoch 才能产生良好结果的技术。
随机擦除 (Random Erasing)
另一种已知有助于提高分类准确率的数据增强技术是随机擦除 (Random Erasing) [10], [11]。它通常与自动增强方法配对使用，由于其正则化效果，通常会带来准确率的额外提升。在我们的实验中，我们仅通过网格搜索调整了应用该方法的概率，发现将其概率保持在低水平（通常约为 10%）是有益的。
以下是在之前基础上引入的额外参数：
random_erase=0.1,
应用随机擦除使我们的 Acc@1 进一步提高了 0.190 个点。
标签平滑 (Label Smoothing)
减少过拟合的一种好方法是防止模型变得过于自信。这可以通过使用标签平滑 (Label Smoothing) [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 个点。
权重衰减 (Weight Decay) 调整
我们的标准配方使用 L2 正则化来减少过拟合。权重衰减参数控制正则化的程度（越大越强），默认情况下普遍应用于模型的所有学习参数。在此配方中，我们对标准方法应用了两次优化。首先，我们执行网格搜索来调整权重衰减参数；其次，我们对归一化层的参数禁用权重衰减。
下面您可以找到我们配方中权重衰减的最佳配置：
weight_decay=2e-05,
norm_weight_decay=0.0,
上述更新使我们的准确率进一步提高了 0.526 个点，为调整权重衰减对模型性能有显著影响这一已知事实提供了额外的实验证据。我们分离归一化参数的方法受到了 ClassyVision 方法的启发。
FixRes 缓解措施
我们在实验初期发现的一个重要特性是，如果验证时使用的分辨率比训练时的 224×224 提高，模型的表现会显著更好。这一效应在 FixRes 论文 [5] 中得到了详细研究，并提出了两种缓解措施：a) 尝试降低训练分辨率，以使验证分辨率下的准确率最大化；或者 b) 在两阶段训练中对模型进行微调，以使其适应目标分辨率。由于我们不想引入两阶段训练，我们选择了选项 a)。这意味着我们减小了训练裁剪尺寸 (train crop size) 并使用网格搜索来寻找在 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 类构建的，关键区别在于它不仅平均模型参数，还平均其缓冲区 (buffers)。此外，我们采用了来自 Pycls 的技巧，使我们能够以一种不依赖于 epoch 数的方式对衰减进行参数化。
推理重采样调整
与过程中所有其他涉及使用不同参数训练模型的步骤不同，此优化是在最终模型的基础上完成的。在推理过程中，图像被缩放到特定分辨率，然后从中提取一个中心的 224×224 裁剪。原始配方使用 256 的重采样尺寸，这导致了与 FixRes 论文 [5] 中描述的类似的差异。通过使此重采样值更接近目标推理分辨率，可以提高准确率。为了选择该值，我们在 [224, 256] 区间内以 8 为步长运行了短网格搜索。为了避免过拟合，该值是使用一半验证集选择的，并使用另一半进行验证。
下面您可以看到我们配方中使用的最佳值：
val_resize_size=232,
以上是使我们的准确率提高 0.224 个点的优化。值得注意的是，ResNet50 的最佳值对于 ResNet101、ResNet152 和 ResNeXt50 也同样适用，这暗示它在不同模型间具有通用性。
[更新] 重复增强 (Repeated Augmentation)
重复增强 (Repeated Augmentation) [15], [16] 是另一种可以提高整体准确率的技术，已被其他强力配方（如 [6], [7] 中的那些）使用。社区贡献者 Tal Ben-Nun 进一步改进了我们最初的配方，提出了用 4 次重复来训练模型。他的贡献是在本文发布后提出的。
下面您可以看到我们配方中使用的最佳值：
ra_sampler=True,
ra_reps=4,
以上是使我们的准确率提高 0.184 个点的最后一次优化。
测试过但未采用的优化
在我们研究的早期阶段，我们尝试了额外的技术、配置和优化。由于我们的目标是保持配方尽可能简单，我们决定不包含任何没有带来显著提升的内容。以下是我们尝试过但没有进入最终配方的几种方法：
- 优化器： 使用更复杂的优化器（如 Adam、RMSProp 或带有 Nesterov 动量的 SGD）并没有比带有动量的原生 SGD 提供明显更好的结果。
- LR 调度器： 我们尝试了不同的 LR 调度方案，如 StepLR 和 Exponential。虽然后者倾向于与 EMA 配合得更好，但它通常需要额外的超参数（例如定义最小 LR）才能工作良好。相反，我们只是使用余弦退火，将 LR 衰减到零，并选择准确率最高的检查点。
- 自动增强： 我们尝试了不同的增强策略，如 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] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. “Deep Residual Learning for Image Recognition”.
- [2] Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, Mu Li. “Bag of Tricks for Image Classification with Convolutional Neural Networks”
- [3] Piotr Dollár, Mannat Singh, Ross Girshick. “Fast and Accurate Model Scaling”
- [4] Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, Ross Girshick. “Early Convolutions Help Transformers See Better”
- [5] Hugo Touvron, Andrea Vedaldi, Matthijs Douze, Hervé Jégou. “Fixing the train-test resolution discrepancy
- [6] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou. “Training data-efficient image transformers & distillation through attention”
- [7] Ross Wightman, Hugo Touvron, Hervé Jégou. “ResNet strikes back: An improved training procedure in timm”
- [8] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, Vaishaal Shankar. “Do ImageNet Classifiers Generalize to ImageNet?”
- [9] Samuel G. Müller, Frank Hutter. “TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation”
- [10] Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, Yi Yang. “Random Erasing Data Augmentation”
- [11] Terrance DeVries, Graham W. Taylor. “Improved Regularization of Convolutional Neural Networks with Cutout”
- [12] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, Zbigniew Wojna. “Rethinking the Inception Architecture for Computer Vision”
- [13] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz. “mixup: Beyond Empirical Risk Minimization”
- [14] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, Youngjoon Yoo. “CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features”
- [15] Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, Daniel Soudry. “Augment your batch: better training with larger batches”
- [16] Maxim Berman, Hervé Jégou, Andrea Vedaldi, Iasonas Kokkinos, Matthijs Douze. “Multigrain: a unified image embedding for classes and instances”