由 Jessica Yung 和 Joan Puigcerver 发布

在本文中，我们将向您介绍如何使用 BigTransfer (BiT)，这是一组预训练的图像模型，可以将其迁移到新数据集以获得出色的性能，即使每个类别只有几个示例。

在 ImageNet 上预训练的 ResNet50 是目前行业中提取图像表示的标准方法。在我们的 BigTransfer (BiT) 论文 中，我们分享了在许多任务中表现明显更好的模型，即使在每个数据集仅使用少数图像的情况下也能很好地迁移。

您可以在 TFHub 中找到在 ImageNet 和 ImageNet-21k 上预训练的 BiT 模型，它们以 TensorFlow2 SavedModels 的形式保存，您可以轻松地将其用作 Keras 层。模型有各种尺寸，从标准 ResNet50 到 ResNet152x4（152 层深，比典型 ResNet50 宽 4 倍），适用于拥有更大计算和内存预算但需要更高精度要求的用户。

图 1：x 轴显示每个类别使用的图像数量，范围从 1 到完整数据集。在左侧的图表中，上面的蓝色曲线是我们的 BiT-L 模型，而下面的曲线是在 ImageNet（ILSVRC-2012）上预训练的 ResNet-50。

在本教程中，我们将展示如何加载我们的 BiT 模型之一，并将其用于以下两种方式：（1）开箱即用，或者（2）微调以适应您的目标任务，以获得更高的精度。具体来说，我们将演示使用在 ImageNet-21k 上训练的 ResNet50。

什么是 Big Transfer (BiT)？

在我们深入了解如何使用这些模型之前，我们是如何训练可以很好地迁移到许多任务的模型的？

上游训练

其本质就在于它的名字 - 我们有效地在大数据集上训练大型架构。在我们的论文发表之前，很少有论文发现通过在大规模公共数据集（如 ImageNet-21k（1400 万张图像，比常用的 ImageNet 大 10 倍））上训练，能够获得显著的收益。我们提炼出的用于训练可以很好地迁移模型的组件包括

大数据集
我们所有模型中最好的性能会随着数据集大小的增加而提高。

大型架构
我们证明，为了充分利用大数据集，需要足够大的架构。例如，在 JFT（3 亿张图像）上训练 ResNet50 并不总是比在 ImageNet-21k（1480 万张图像）上训练 ResNet50 性能更好，但当我们在 JFT 上训练更大的模型（例如 ResNet152x4）时，我们始终看到性能提升，与在 ImageNet-21k 上训练相比（参见下面的图 2）。

图 2：上游数据集（x 轴）和模型大小（气泡大小/颜色）对下游任务性能的影响。仅使用更大的数据集或更大的模型可能会损害性能 - 两种都需要同时增加。

较长的预训练时间
我们还证明，在大型数据集上预训练时，训练时间足够长也很重要。在 ImageNet 上训练 90 个 epochs 是标准做法，但如果我们在更大的数据集（如 ImageNet-21k）上训练相同的步数（然后在 ImageNet 上微调），那么性能会比直接在 ImageNet 上训练差。

GroupNorm 和权重标准化
最后，我们使用 GroupNorm 结合权重标准化，而不是 BatchNorm。由于我们的模型很大，我们只能在每个加速器（例如 GPU 或 TPU 芯片）上拟合少数图像。但是，当每个加速器上的图像数量太少时，BatchNorm 的性能会下降。GroupNorm 没有这个问题，但不能很好地扩展到大型全局批处理大小。但是，当我们将 GroupNom 与权重标准化结合起来时，我们发现 GroupNorm 能够很好地扩展到大型批处理大小，甚至比 BatchNorm 性能更好。

下游微调

此外，下游微调在数据效率和计算量方面成本低廉 - 我们的模型在自然图像上仅使用少数示例就能获得良好的性能。我们还设计了一种超参数配置，我们称之为“BiT-HyperRule”，它在许多任务上表现相当好，无需进行昂贵的超参数搜索。

BiT-HyperRule：我们的超参数启发式方法
如上所述，这不是超参数搜索 - 给定一个数据集，它指定了一组我们发现可以产生良好结果的超参数。通过运行更昂贵的超参数搜索，您通常可以获得更好的结果，但 BiT-HyperRule 是在您的数据集上获得良好初始结果的有效方法。

在 BiT-HyperRule 中，我们使用 SGD，初始学习率为 0.003，动量为 0.9，批处理大小为 512。在微调期间，我们在训练步骤的 30%、60% 和 90% 时将学习率衰减 10 倍。

作为数据预处理，我们调整图像大小，进行随机裁剪，然后进行随机水平翻转（详细信息见表 1）。我们对所有任务都进行随机裁剪和水平翻转，除了那些会导致标签语义发生破坏的任务。例如，我们不会对计数任务应用随机裁剪，也不会对我们应该预测对象方向的任务应用随机水平翻转（见图 3）。

表 1：下游调整大小和随机裁剪详细信息。如果图像更大，我们将它们调整为更大的固定大小，以从更高分辨率的微调中获益。

图 3：CLEVR 计数示例：这里任务是统计图像中小型圆柱体或红色物体的数量。我们不会应用随机裁剪，因为这可能会裁剪掉我们想要计数的物体，但我们会应用随机水平翻转，因为它不会改变我们关心的图像中物体数量（因此不会改变标签）。图像归属：Johnson 等人的 CLEVR 计数示例)

我们根据数据集大小（见表 2）确定时间表长度以及是否使用 MixUp（Zhang 等人，2018 年，图 4 说明了这一点）。

图 4：MixUp 接受示例对，并线性组合图像和标签。这些图像取自 tf_flowers 数据集。

表 2：关于下游时间表长度和我们何时使用 MixUp 的详细信息。

我们根据经验结果确定了这些超参数启发式方法。我们在 论文 和我们的 Google AI 博客文章 中详细解释了我们的方法，并描述了我们的结果。

教程

现在让我们实际微调这些模型之一！您可以通过在 此 colab 中运行代码来进行操作。

1）加载预训练的 BiT 模型

您可以从 TensorFlow Hub 下载我们预训练在 ImageNet-21k 上的 BiT 模型之一。这些模型以 SavedModels 的形式保存。加载它们非常简单

import tensorflow_hub as hub
# Load model from TFHub into KerasLayer
model_url = "https://tfhub.dev/google/bit/m-r50x1/1"
module = hub.KerasLayer(model_url)

2）开箱即用 BiT

如果您还没有为您的图像准备标签（或者只是想玩玩），您可能对开箱即用地使用模型感兴趣，即不进行微调。为此，我们将使用在 ImageNet 上微调的模型，以便它拥有可解释的 1000 个类别的 ImageNet 标签空间。许多常见对象没有被涵盖，但它可以让我们大致了解图像中包含什么。

# use model
logits = imagenet_module(image)

请注意，BiT 模型接受值为 0 到 1 之间的输入。

在 colab 中，您可以从 URL 加载图像，并查看模型的预测结果

> show_preds(preds, image[0])

图像来自 PikRepo

在这里，预训练在 ImageNet 上的模型正确地将照片分类为大象。由于耳朵的大小，它更有可能是一只印度大象而不是非洲大象。在 colab 中，我们还对要进行微调的数据集（TF flowers）中的图像进行了预测，该数据集也已用于 其他教程。请注意，正确的标签“郁金香”不是 ImageNet 中的一个类别，因此模型目前无法预测它 - 让我们看看它会尝试做什么

该模型预测了一个看起来相当相似的类别，“甜椒”。

3）在您的任务上微调 BiT

现在，我们将对 BiT 模型进行微调，使其在特定数据集上表现更好。这里我们将使用 Keras 为了简单起见，我们将使用一个花卉数据集（tf_flowers）对模型进行微调。我们将使用在开头加载的模型（即在 ImageNet-21k 上预训练的模型），这样它就不会对狭窄的类别子集产生偏见。

有两步

1. 创建一个具有新最终层（称为“头部”）的新模型
   
2. 使用 BiT-HyperRule（我们的超参数启发式方法）微调此模型。我们在本文的“下游微调”部分详细介绍了这一点。
   

要创建新模型，我们

1. 切断 BiT 模型的原始头部。这将留下“预 logits”输出。
   
   • 如果我们使用“特征提取”模型，则不必执行此操作，因为这些模型的头部已经切断了。
     
2. 添加一个新的头部，其输出数量等于新任务的类别数量。请注意，将头部初始化为全零非常重要。
   

class MyBiTModel(tf.keras.Model):
  """BiT with a new head."""

  def __init__(self, num_classes, module):
    super().__init__()

    self.num_classes = num_classes
    self.head = tf.keras.layers.Dense(num_classes, kernel_initializer='zeros')
    self.bit_model = module

  def call(self, images):
    # No need to cut head off since we are using feature extractor model
    bit_embedding = self.bit_model(images)
    return self.head(bit_embedding)

model = MyBiTModel(num_classes=5, module=module)

在微调模型时，我们使用 BiT-HyperRule，这是我们之前描述的用于选择下游微调超参数的启发式方法。我们还在 colab 中完整地编写了我们的启发式方法。

# Define optimiser and loss
 
# Decay learning rate by factor of 10 at SCHEDULE_BOUNDARIES.
lr = 0.003
SCHEDULE_BOUNDARIES = [200, 300, 400, 500]
lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries=SCHEDULE_BOUNDARIES,
                                                                  values=[lr, lr*0.1, lr*0.001, lr*0.0001])
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.9)

为了微调模型，我们使用简单的 Keras model.fit API

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=optimizer,
             loss=loss_fn,
             metrics=['accuracy'])
 
# Fine-tune model
model.fit(
   pipeline_train,
   batch_size=512,
   steps_per_epoch=10,
   epochs=50,
   validation_data=pipeline_test)

我们看到我们的模型在 20 步内达到了 95% 的验证准确率，并在使用 BiT-HyperRule 微调后达到了超过 98% 的验证准确率。

4) 保存微调后的模型以供日后使用
保存模型以供日后使用非常容易。您可以像我们在开始时加载 BiT 模型一样加载保存的模型。

# Save fine-tuned model as SavedModel
export_module_dir = '/tmp/my_saved_bit_model/'
tf.saved_model.save(model, export_module_dir)
 
# Load saved model
saved_module = hub.KerasLayer(export_module_dir, trainable=True)

瞧！我们现在有一个模型可以预测郁金香而不是甜椒。

摘要

在这篇文章中，您了解了可以用来训练模型的关键组件，这些模型可以很好地转移到许多不同的任务中。您还了解了如何加载我们的一个 BiT 模型，在您的目标任务上对其进行微调并保存结果模型。希望这有所帮助，祝您微调愉快！

致谢
这篇文章基于 Alexander Kolesnikov、Lucas Beyer、Xiaohua Zhai、Joan Puigcerver、Jessica Yung、Sylvain Gelly 和 Neil Houlsby 的工作。我们感谢苏黎世大脑研究中心和 TensorFlow 团队的许多成员的反馈，特别是 Luiz Gustavo Martins、André Susano Pinto、Marcin Michalski、Josh Gordon、Martin Wicke、Daniel Keysers、Amélie Royer、Basil Mustafa 和 Mario Lučić。

其他链接

• 教程 Colab
  
• TensorFlow Hub 上的模型
  
• GitHub 仓库
  
• BigTransfer (BiT) 论文
  
• BiT Google AI 博客文章