开源 BiT：计算机视觉大规模预训练的探索

2020 年 7 月 6 日 谷歌开发者

文 / Lucas Beyer 和 Alexander Kolesnikov，研究工程师，Google Research

计算机视觉领域的研究者普遍认为，现代深度神经网络的发展始终需要使用大量标记数据，如，当前最前沿 (SOTA) 的卷积神经网络 (CNN) 便需要在 OpenImages 或 Places（其中包含超过 100 万张标签图像）等数据集上进行训练。但是，对很多普通的应用开发者来说，收集到足够多的标记数据颇为困难。

为缓解计算机视觉任务缺乏标记数据的情况，我们通常会使用由通用数据（如 ImageNet） 预训练 的模型，即将通用数据上习得的视觉特征在兴趣任务上重复使用。尽管这种预训练模型在实践中相当有效，但在多环境中快速掌握和新概念理解上，仍存在短板。鉴于 BERT 和 T5 在语言领域取得的进展，我们认为 大规模 预训练同样能提升计算机视觉模型的性能。

T5
https://ai.googleblog.com/2020/02/exploring-transfer-learning-with-t5.html

在《Big Transfer (BiT)：通用视觉表征学习》 (Big Transfer (BiT): General Visual Representation Learning) 一文中，我们设计出一种新方法，利用超出业界标准规模 (ILSVRC-2012) 的图像数据集对一般特征进行有效的预训练。值得一提的是，随着预训练数据量的增加，选择合适的归一化层和适当扩展架构容量十分重要。

Big Transfer (BiT)：通用视觉表征学习
https://arxiv.org/abs/1912.11370

新方法表现出前所未有的性能，可适应各种新的视觉任务，包括小样本识别设置和最近引入的“现实世界 (ObjectNet)” 基准。我们很高兴能分享在公开数据集上预训练的最佳 BiT 模型，以及 TF2、Jax 和 PyTorch 中的代码。借助这些内容，任何人都能在他们感兴趣的任务中获得出色性能，即使每个类只有少量标签图像也是如此。

最佳 BiT 模型
https://tfhub.dev/google/collections/bit/1

预训练

为了研究数据规模对模型的影响，我们通过三个数据集重新审视了预训练设置中常用的设计选项（如激活函数和权重的归一化、模型宽度/深度和训练计划）：ILSVRC-2012（128 万张图像，1000 个类）、ImageNet-21k（1400 万张图像，约 21000 个类）以及 JFT（3 亿张图像，大约 18000 个类）。更重要的是，利用这些数据集，我们能够专心研究先前未曾充分探索的大型数据模式。

我们首先研究了数据集规模与模型容量之间的相互作用。为此，我们选择了经典的 ResNet 架构进行训练。该架构性能良好，同时简单且可重现。从标准 50 层深的“R50x1”到 4 倍宽度、152 层深的“R152x4”的变体模型，我们在上述每个数据集上都进行了训练。我们通过观察发现了一个重要结论：要充分利用大量数据的优势，就必须增加模型容量。下图左侧面板中的红色箭头便证明了这一结论：

左侧：想要有效利用大型数据集进行预训练，就需要增加模型容量。红色箭头证明了以下内容：在较大的 ImageNet-21k 上预训练后，小型架构（较小的点）的性能会变差，而大型架构（较大的点）的性能则有所改善。右侧：仅利用较大数据集进行预训练未必会提高性能，例如，从 ILSVRC-2012 改用相对较大的 ImageNet-21k 后，性能没提升。不过，通过增加计算预算和训练时间，性能明显有所提升

经观察后，我们得出第二个更为重要的结论：训练时间的作用不容忽视。如果不调整计算预算、不增加训练时间，而只增加预训练所用数据集的规模，模型性能可能会变得更差。但是，调整新数据集的训练时间后，性能得到了显著提升。

我们在探索期间发现另一项修改对于性能提升也十分关键。我们由此得出结论，将批标准化（Batch Normalization，BN，一种可通过标准化激活函数实现稳定训练的通用层）替换为组标准化 (Group Normalization，GN) 对于大规模预训练十分有利。

首先，BN 的状态（神经激活函数的均值和方差）需要在预训练和迁移之间进行调整，而 GN 是无状态的，从而避免了这一难题。其次，BN 使用批处理级别的统计数据，对于每个设备上的小型批次大小来说并不可靠，而大型模型无法避免在多设备上进行训练。由于 GN 不需要计算批处理级别的统计信息，因此也避免了这个问题。

请参阅我们的论文，以便了解更多技术细节，包括如何使用权重标准化 (Weight Standardization) 技术来确保行为的稳定。

权重标准化
https://arxiv.org/abs/1903.10520

预训练策略汇总：使用标准的 ResNet，增加深度和宽度，将 BatchNorm (BN) 替换为 GroupNorm 和权重标准化 (GNWS)，然后通过一个大型通用数据集进行训练，并完成多次迭代

迁移学习

根据 BERT 在语言领域建立的方法，我们利用来自感兴趣的各项“下游”任务的数据对预训练的 BiT 模型进行了微调，这些任务可能只有少量标记数据。由于预训练模型对视觉世界的理解能力较高，因此这个简单的策略非常有效。

微调时需要选择大量超参数，如学习率、权重衰减等。我们提出了一个启发式的超参数选择方法，名为 “BiT-HyperRule”。该方法只以高层次的数据集特征为基础，如图像分辨率和标签样本数量。我们已成功将 BiT-HyperRule 方法应用到 20 多项从自然图像到医学影像等不同的任务上。

经预训练后，BiT 模型可根据任何任务进行微调，即使提供的标签样本非常有限也无妨

将 BiT 迁移到仅有少量示例的任务后，我们发现：如果同时增加用于预训练的通用数据量和架构容量，生成的模型明显能更快适应新数据。在拥有 1 个样本和 5 个样本的 CIFAR（见下图）上，如果在 ILSVRC 上进行预训练，则增加模型容量所带来的增益比较有限（绿色曲线）。但是，如果在 JFT 上进行大规模预训练，每次增加模型容量都会带来巨大增益（棕色曲线）。BiT-L 在 1 个样本上的准确率可达 64%，在 5 个样本上的准确率可达 95%。

曲线描绘了超过 5 次独立运行（光点）后，模型迁移到 CIFAR-10（每个类有 1 张或 5 张图像，共 10 张或 50 张图像）上的平均准确率。事实证明，在大型数据集上进行预训练的大型架构更能发挥数据的作用

为了验证该结果具有普遍性，我们还在 VTAB-1k 上对 BiT 进行了评估。VTAB-1k 是由 19 个不同任务组成的任务包，每个任务只有 1000 个标签样本。我们将 BiT-L 模型迁移到所有这些任务上，总体得分为 76.3%，比之前的最高水准高出了 5.8%。

我们根据几个标准计算机视觉基准（如 Oxford Pets、Flowers、CIFAR 等）对 BiT-L 进行了评估，并由此证明，大规模预训练和简单迁移的策略十分有效，即使提供的数据量仅属于中等水平也是如此。综上所述，BiT-L 已达到或超越了目前的最高水准。最后，我们在 MSCOCO-2017 检测任务中将 BiT 用作 RetinaNet 的主干，同时证实，即使是处理这种结构化输出任务，使用大规模预训练也大有裨益。

左侧：在各种标准计算机视觉基准，BiT-L 与先前最前沿的通用模型的准确率。右侧：在 MSCOCO-2017 中将 BiT 用作 RetinaNet 主干所得出的平均精度 (AP) 结果

需要强调的是，在我们考虑的所有不同下游任务中，我们不会对任务逐个进行超参数调整，而是借助 BiT-HyperRule。如我们的论文所述，在足够大的验证数据上调整超参数可以获得更好的结果。

在 ObjectNet 上的评估

为进一步评估 BiT 在更具挑战的场景中的稳健性，我们在最近引入的 ObjectNet 数据集上对经 ILSVRC-2012 微调过的 BiT 模型进行了评估。ObjectNet 更为接近真实场景，其中，对象可能出现在非典型的上下文中，并以非典型的视角和旋转呈现。有趣的是，受益于更大的数据和架构容量，BiT-L 在这个任务上取得了令人瞩目的成就：实现了高达 80.0% 的 top-5 准确率，比之前的最高水准 (SOTA) 高出了近 25%。

ObjectNet
http://objectnet.dev/

BiT 在 ObjectNet 数据集上得出的评估结果。左侧：top-5 准确率，右侧：top-1 准确率

结论

我们发现，在利用大量通用数据进行预训练的情况下，一个简单的迁移策略就会让结果大幅改观，不管是对大型数据集，还是对数据极少、甚至每个类仅有一张图像的任务来说均是如此。

我们已经发布 BiT-M 模型（经 ImageNet-21k 预训练的 R152x4 模型），同时还提供了在 Jax、TensorFlow 2 和 PyTorch 中进行迁移的 Colab。除了发布代码外，我们还向读者介绍了有关如何使用 BiT 模型的 TensorFlow 2 动手教程。我们期待从业者和研究人员可以从中获益，并用该模型替代常用的 ImageNet 预训练模型。

已经发布 BiT-M 模型
https://github.com/google-research/big_transfer

致谢

我们要感谢 Xiaohua Zhai、Joan Puigcerver、Jessica Yung、Sylvain Gelly 和 Neil Houlsby，他们共同撰写了 BiT 论文，并参与了 BiT 开发的各项环节，同时也要感谢苏黎世的 Brain 团队。我们还要感谢 Andrei Giurgiu 在调试输入流水线方面给予的帮助。感谢 Tom Small，是他制作了这篇博文中的动画。最后，感兴趣的读者可以了解一下由 Google Research 的同事开发的一些与本方向相关的方法、Noisy Student 模型，以及 Facebook Research 发布的一篇高度相关的论文探索弱监督预训练的限制 (Exploring the Limits of Weakly Supervised Pretraining)。