一种基于迁移学习的视觉多任务模型探析

刘夏鸣

（华中光电技术研究所-武汉光电国家研究中心，湖北武汉 430223）

1 概述

近年来深度迁移学习在包括计算机视觉、自然语言处理等人工智能相关领域都取得了很大进展。例如基于海量语料的预训练的语言表征模型（Bidirectional Encoder Representation from Transformers，BERT）模型[1]已经被证明可以显著提升几乎所有的自然语言处理任务。在计算机视觉领域，基于上百万图像数据集[2]标签数据训练的分类模型可以迁移到语义分割、目标检测等任务上，大幅提升这些任务的准确率等指标。本文研究的问题是如何在一定的约束条件下用迁移学习的方法同时提升多个任务的指标。我们主要考虑两个在实际应用中广泛存在的约束条件：一，计算资源和存储资源受限。此约束条件在某些硬件例如低功耗设备上格外突出。二，任务之间的耦合度应当尽可能，使得每个任务可以在不影响其它任务的情况下得到迭代更新。

在传统的迁移学习方案中，下游任务模型是在预训练模型的基础上通过“微调”的方式获得的[3]。在微调的过程中，任务模型以预训练模型的网络权重作为初始权重，在任务标注数据上以梯度下降算法进行训练，逐步改变网络权重，最终得到了一个新的模型。当有多个任务存在时，人们需要为每个任务微调得到一个独一无二的任务模型，且这些任务模型中的权重和计算无法共用。显而易见，传统迁移学习方案计算复杂度和存储复杂度都为O(n)，其中n 是任务的数量。这样的复杂度在大多数应用中是不可接受的。

2 模型设计

本文考虑拥有12 个Transformer 模块的ViT-Base 模型。我们提出的方案流程由三个步骤构成：（1）单任务微调训练；（2）单任务模型蒸馏；（3）多任务模型合并。模型方案的细节如下：

2.1 主干模型

2.1.1 Transformer 编码器

本文使用的ViT-Base 模型拥有12 层Transformer 编码器网络。Transformer 模型在2017 年由A. Vaswani 等人提出，原用于机器翻译任务。目前基于Transformer 结构的模型已经在许多自然语言处理和计算机视觉任务中取得了远超传统卷积神经网络（CNN）或长短时期记忆网络（LSTM）的表现。Transformer 网络有两种类型，分别是编码器和解码器。一个Transformer 编码器模块由若干个编码器和解码器堆叠形成。如下图的左侧部分为编码器，由多头注意（Multi-Head Attention）和一个全连接组成，用于将输入语料转化成特征向量。右侧部分是解码器，其输入为编码器的输出以及已经预测的结果，保证训练与推理的一致性，用于输出最后结果的条件概率，图1。

图1 Transformer 结构图

2.1.2 ViT 模型结构

在计算机视觉有关的任务中，要让Transformer 模型能够接受图像输入，我们必须设法将一副二维图像转化成一个一维的向量序列。在ViT 模型中，我们把一副高度为H 宽度为W 的图像拆分成N 个边长为P 的正方形图像小块，其中N=HW/P2。若图像中的每个像素包含C 个通道，则每个正方形小块对应一个维度为CP2 的向量，该向量记录了图像小块的每个像素每个通道的原始数值。ViT 模型用一个可学习的全连接网络将该向量非线性映射成了一个D 维的特征向量。通过这种方法，原始图像变成了一个包含N 个D 维向量的序列，后者可以作为Transformer层的输入。为了区分每个D 维向量对应的正方形小块在原始图像中的位置，ViT 为每个位置的小块引入了一个可学习的位置向量。该位置向量会加到特征向量中，作为后续Transformer 层的输入。最后，ViT 引入了一个类似BERT 的[CLS]向量作为零号特征向量，该向量作为整幅图像的特征表示。

2.1.3 ViT 模型预训练

我们使用的ViT-Base 模型在图像网络数据集上进行预训练。图像网络数据集是一个图像分类数据集，其训练集、验证集和测试集分别包含128 万、5 万和10 万张标注图片，分属于1000 个类别。我们使用的ViT-Base模型在预训练任务上经过300 次循环后最终可以达到84.15%的准确率。

2.2 多任务模型训练流程

2.2.1 单任务模型微调

第一个步骤，是针对每个任务微调预训练ViT 模型深层（靠近输出层）的L 个模块，冻结剩余位于浅层的12-L 个模块。具体的L 数值是一个超参数，可以随着任务的不同而变化。一个任务拥有较大的L 值意味着该任务被分配了更多的可训练参数以及计算资源。在极端情况下L=12，此时我们的微调方案和传统的全微调等价。为了确保最低程度的参数和计算共享，我们在一个预先设定好的较小的区间N_min <= L <= N_max 内搜寻最佳的L 值。在我们后续的实验中，我们取N_min=4 以及N_max=10。

第一个步骤最终为每一个任务输出一个任务模型，这些任务模型的底层和预训练模型共享参数，顶层使用任务独占的网络参数。我们将这些模型称为教师模型（见图2 左），用于后续的模型蒸馏。

图2 多任务ViT 模型设计图

2.2.2 单任务模型蒸馏

在上一个步骤中，我们用微调部分网络参数的办法针对每个任务微调得到了一个任务模型。这些任务模型没有直接的联系，因此可以用知识蒸馏的方法压缩这些模型的非共享网络。具体的说，对一个第一步中得到的有着L_t 层任务独占网络以及12-L_t 层冻结网络任务模型，我们可以用任意一种模型蒸馏算法将前者压缩成一个仅有l_t 层的更小的模块，而让后者保持不变(见图2 中）。在实验中，我们发现让学生模型的网络参数从教师模型的对应网络层中初始化可以达到最佳的学习效果。

2.2.3 多任务模型合并

在最后一步里，我们把所有的单任务学生模型合并成一个多任务模型，这样冻结网络层中的参数和所有计算可以被各个任务所共享，从而提升效率（见图2 右）。在实际操作中，我们只需要在创建多任务模型计算图后将各任务模型中的参数载入该计算图中的对应位置即可。

3 试验方案

3.1 基线模型

在此我们将对比在一系列任务上我们的模型与主流基线模型的表现。

3.1.1 单任务模型（无模型蒸馏）：此模型由预训练模型ViT 在任务数据上以全参数微调的方式得到。此为预训练模型进行迁移学习的最佳实践。如图3 所示。

图3 预训练ViT 模型结构图

3.1.2 单任务模型（有模型蒸馏）我们把拥有12 个Transformer 层的单任务模型通过模型蒸馏的方式压缩至6 个模块，从而降低计算复杂度。

3.1.3 多任务模型：采用传统的多任务学习的方式训练，除了各任务输出层以外共享全部的12 层传递参数，具体方案为每个任务分配任务专属的1～3 个传递层用于学习任务相关的特征。所有任务共享预训练模型的前6 层网络作为特征提取层。为了更加清晰的展示ViT 模型结构和训练过程中的向量变化，图4 给出了ViT的向量维度变化图。

图4 ViT 的向量维度变化图

3.2 任务与数据集

3.2.1 分类任务：分类样本包含100 个类别，每个类别共600 张图片作为训练样本，每张图片的分辨率为32x32，如图5 示例。

图5 分类样本示例

3.2.2 目标检测：目标样本包含20 万张图片，50 万标注目标，共80 个类别。

3.2.3 语义分割：本样本包含5000 张细粒度标注的图片，2 万张粗粒度标注的图片，标注共分属于30 个类别。

3.3 参数选择与试验

在接下来的实验中使用Adam 的优化器，采用学习率为2e-5 的32 批次训练超参数，从实验结果中可以得出，我们的模型在图像分类、语义分割和目标检测这三个经典的视觉任务上都超过了单任务模型蒸馏和多任务联合训练的基线结果，仅仅只是略微逊于单任务模型全微调时的表现。仅使用3 层传递网络，我们在分类任务、目标检测和语义分割这三个数据集上任务上分别达到了12 层基线网络指标的99.4%，98.3%，98.5%，如表1 所示。

表1 模型在三个视觉任务上的表现对比

4 模型分析

4.1 优势

本文提出的模型方案的主要优点是它的灵活性。首先，不同的视觉任务可以从不同深度的网络层获取不同抽象层级的表征。其次，我们的方案可以允许给不同的任务分配不同的计算资源，例如更重要的任务可以分配更多的计算资源，或者降低资源投入边际效费比较低的任务资源，从而使得模型在多任务环境下的平均表现打到更优。第三，通过对任务之间的依赖进行解耦，我们的模型有效回避了传统多任务模型的任务之间相互冲突的可能性，也不再需要针对任务权重超参数进行细致的调整。我们的多任务模型可以非常直观的对单个任务子网络模型的参数进行更新。

4.2 与适配器模型的比较

适配器模型是另一类在资源受限条件下应用较为广泛的模型。适配器模型通过在一个模型网络的层与层之间添加一个小型可训练模块（即适配器）来实现预训练模型对特定任务的适配。在迁移学习的过程中，原预训练模型中既有参数不发生变化，仅有适配器中的参数得到更新。这类方法的优势是只需要引入极少量的参数即可实现较优秀的任务表现，但其劣势在于计算开销不会降低，因为任务网络中的数据流被适配器完全改变了，无法跨任务复用。也就是说，在多任务场景下适配器模型的计算复杂度仍为O(n)，因此无法与本文中提出的方法进行直接的比较。

结束语

本文提出了一种新型多任务模型框架，可以在保证灵活性、低耦合性的同时极大降低模型对内存的开销和对计算资源的开销。我们在多个视觉任务上证明了我们可以在保证模型预测精度的同时降低至多75%的资源开销，因此具有较强的应用推广价值。