BigTransfer (BiT)：计算机视觉领域最前沿迁移学习模型

磐创AI分享

来源 | TensorFlow

作者 | Jessica Yung 、Joan Puigcerver

我们将在本文中为您介绍如何使用 BigTransfer (BiT)。BiT 是一组预训练的图像模型：即便每个类只有少量样本，经迁移后也能够在新数据集上实现出色的性能。

经 ImageNet 预训练的 ResNet50 系列模型是当今图像提取表征的业界标准，而我们在 BigTransfer (BiT) 论文中分享的模型在跨多任务上的性能明显优于 ResNet50，即便每个数据集只使用少数几张图像，也能取得不俗的迁移效果。

• BigTransfer (BiT) 论文 https://arxiv.org/abs/1912.11370

您可以在 TFHub 中找到在 ImageNet 和 ImageNet-21k 上预训练的 BiT 模型，并像使用 Keras 层一样，轻松使用 TensorFlow2 SavedModel。这些模型有各种规模，包括标准的 ResNet50 和 ResNet152x4（152 层深，比典型的 ResNet50 宽 4 倍），后者适合计算和内存预算大、对准确率要求更高的用户。

• TFHub https://tfhub.dev/google/collections/bit/1

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

在本教程中，我们将展示如何加载其中一种 BiT 模型，并：

1. 以原生方式使用模型或
2. 针对目标任务微调模型以提高准确率

具体来说，我们将演示如何使用在基于 ImageNet-21k 上训练的 ResNet50。

什么是 BigTransfer (BiT)？

在了解模型的详细使用方法之前，我们首先要了解如何训练此类模型，使其可有效迁移至多个任务。

上游训练

上游训练的精髓就体现在其名称，即我们可以在大数据集上有效地训练大型架构。在我们的论文发表前，很少有在 ImageNet-21k（拥有 1400 万张图像，比常用的 ImageNet 大 10 倍）等大型公开数据集上取得显著训练效果的论文公布。为了训练可实现有效迁移的模型，我们精选出以下组件：

大型数据集

随着数据集规模的增加，模型的最佳性能也会随之提升。

大型架构

我们发现，如果要充分利用大数据集，就需要有足够大的架构。例如，相对于在 ImageNet-21k（拥有 1480 万张图像）上进行 ResNet50 的训练，在 JFT（拥有 3 亿张图像）上训练 ResNet50 未必总能获得性能提升；但在训练 ResNet152x4 等大型模型时，在 JFT 上的性能始终要高于 ImageNet-21k 上的性能（如下方图 2 所示）。

图 2：大型上游数据集（x 轴）和模型大小（气泡大小/颜色）对下游任务性能的影响：单独使大型数据集或模型可能会有损性能，因此二者需要同步增加

足够的预训练时间

我们还发现，在大型数据集上进行预训练时，训练时间也很重要。标准做法是在 ImageNet 上训练 90 个周期。但是，如果在 ImageNet-21k 等大型数据集上进行步数相同的训练（然后在 ImageNet 上进行微调），其性能会比直接在 ImageNet 上训练要差。

组归一化和权重标准化

最后，我们要将组归一化（GroupNorm，而非 BatchNorm）与权重标准化结合使用。由于模型巨大，我们只能在每个加速器（如 GPU 或 TPU 芯片）上拟合几张图像。但当每个加速器上的图像数量过少时，BatchNorm 的性能就会变差。虽然 GroupNorm 没有这个问题，但也无法很好地扩展至整个的大型批次大小。不过，将 GroupNom 与权重标准化结合使用后，我们发现 GroupNorm 表现出针对大型批次大小的良好扩展性，其表现甚至优于 BatchNorm。

下游微调

就数据效率和计算成本而言，每类自然图像仅需少量样本，我们的模型就能获得出色性能，下游微调成本较低。此外，我们还设计了一个名为“BiT-HyperRule”的超参数配置，该配置在许多任务中均表现出色，同时无需进行昂贵的超参数扫描分析。

BiT-HyperRule：超参数启发式配置

如上文所述，此配置无需进行超参数扫描分析：给定数据集后，此配置就会指定一组经证实可取得良好结果的超参数。运行成本高的超参数扫描分析自然能够取得更好的结果，但 BiT-HyperRule 这种有效的方法也可在数据集上取得良好的初始结果。

在 BiT-HyperRule 中，我们采用了初始学习率 0.003、动量 0.9，批次大小 512 的随机梯度下降法 (SGD)。在微调过程中，我们依次在 30%、60% 和 90% 的迭代中将学习率衰减 1/10。

在数据预处理过程中，我们调整了图像大小、随机裁剪，并进行随机水平翻转（详情请参见表 1）。除了标签语义会受随机裁剪和水平翻转操作破坏的任务，我们对所有任务都执行了这类操作。举例来说，我们不会对计数任务进行随机裁剪，也不会对旨在预测物体方向的任务进行随机水平翻转（图 3）。

表 1：下游大小调整和随机裁剪详情。如果图像较大，我们会将其调整到更大的固定尺寸，以便在更高分辨率上更好地进行微调

图 3：CLEVR 计数示例：这里的任务是统计图像中的小圆柱体或红色物体的数量。我们不会进行随机裁剪，因为这类操作可能会裁剪掉需要统计的物体；但我们会进行随机水平翻转，因为这不会改变图像中需要统计对象的数量（因此也不会改变标签）图像提供方：CLEVR 计数示例图像由 Johnson 等人提供

我们需要根据数据集大小（表 2）来确定计划时长，以及是否使用 MixUp（由 Zhang 等人于 2018 年提供，如图 4 所示）。

图 4：MixUp 采用成对样本，并对图像和标签进行了线性组合。这些图像均取自数据集 tf_flowers

表 2：下游计划时长及何时使用 MixUp 的详情

我们根据实验结果确定了这些超参数启发式配置，并在论文和 Google AI 博文中详细介绍了采用的方法和取得的结果。

教程

现在，让我们正式开始微调上文提到的其中一种模型！您可以运行此 Colab 中的代码，跟着我们逐步进行操作。

• 此 Colab https://colab.research.google.com/github/google-research/big_transfer/blob/master/colabs/big_transfer_tf2.ipynb

1) 加载预训练的 BiT 模型

您可以访问 TensorFlow Hub，下载基于 ImageNet-21k 预训练的其中一种 BiT 模型。这类模型会保存为 SavedModel。加载过程非常简单：

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)

• SavedModel https://tensorflow.google.cn/hub/tf2_saved_model

2) 以原生方式使用 BiT 模型

如果您尚未对图像添加标签（或只想随便试试），则可以尝试以原生方式使用模型（即无需进行微调）。为此，我们将使用在 ImageNet 上完成微调的模型，这样便拥有可解释的 1000 个类的 ImageNet 标签空间。模型并未涵盖许多常见对象，但却合理解释了图像内容。

# use model
logits = imagenet_module(image)

请注意：BiT 模型的输入值介于 0 到 1 之间。

• 1000 个类的 ImageNet 标签空间 https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a

在 Colab 中，您可以通过网址加载图像并查看模型的预测结果：

> show_preds(preds, image[0])

图像来源：PikRepo

经 ImageNet 预训练的模型可以将照片正确归类为大象。同时，鉴于耳朵的尺寸，这头大象更有可能是印度象，而不是非洲象。在 Colab 中，我们还对需要微调 tf_flowers数据集中的图像以进行了预测。其他教程中同样也使用了此数据集。请注意，正确的标签“郁金香”并非 ImageNet 中的类，因此模型目前无法进行预测。让我们看看模型会将图像归入哪个类：

该模型预测了一个相似度非常高的类：“灯笼椒”。

• tf_flowers https://tensorflow.google.cn/datasets/catalog/tf_flowers
• 其他教程 https://tensorflow.google.cn/tutorials/load_data/images

3) 针对任务对 BiT 进行微调

现在，我们将对 BiT 模型进行微调，以改善其在特定数据集上的性能。简单起见，我们会使用 Keras，同时将在花朵数据集 (tf_flowers) 上对模型进行微调。我们将使用最初加载的模型（即在 ImageNet-21k 上完成预训练的模型），以免过度偏向各类的一小部分子集。

以下是需要执行的两个步骤：

1. 创建具有全新最终层（称为“头部”）的新模型
2. 使用超参数启发式配置 BiT-HyperRule 对模型进行微调。我们早前已在本文的“下游微调”部分对此进行了详细说明。

要创建新模型，我们需要：

1. 截断 BiT 模型的原始头部，从而获得“pre-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)

为了微调模型，我们使用简单的model.fitAPI：

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 的研究成果为基础。此外，我们还要感谢提供苏黎世 Brain Research 团队和 TensorFlow 团队的众多成员提供的反馈，特别鸣谢 Luiz Gustavo Martins、André Susano Pinto、Marcin Michalski、Josh Gordon、Martin Wicke、Daniel Keysers、Amélie Royer、Basil Mustafa 以及 Mario Lučić。

其他链接

• Colab 教程 https://colab.research.google.com/github/google-research/big_transfer/blob/master/colabs/big_transfer_tf2.ipynb
• TensorFlow Hub 上的模型 https://tfhub.dev/google/collections/bit/1
• GitHub 代码库 https://github.com/google-research/big_transfer
• BigTransfer (BiT) 论文 https://arxiv.org/abs/1912.11370
• BiT Google AI 博文 https://ai.googleblog.com/2020/05/open-sourcing-bit-exploring-large-scale.html

- End -