一種基于遷移學(xué)習(xí)的視覺多任務(wù)模型探析

劉夏鳴

（華中光電技術(shù)研究所-武漢光電國家研究中心，湖北武漢 430223）

1 概述

近年來深度遷移學(xué)習(xí)在包括計算機視覺、自然語言處理等人工智能相關(guān)領(lǐng)域都取得了很大進展。例如基于海量語料的預(yù)訓(xùn)練的語言表征模型（Bidirectional Encoder Representation from Transformers，BERT）模型[1]已經(jīng)被證明可以顯著提升幾乎所有的自然語言處理任務(wù)。在計算機視覺領(lǐng)域，基于上百萬圖像數(shù)據(jù)集[2]標簽數(shù)據(jù)訓(xùn)練的分類模型可以遷移到語義分割、目標檢測等任務(wù)上，大幅提升這些任務(wù)的準確率等指標。本文研究的問題是如何在一定的約束條件下用遷移學(xué)習(xí)的方法同時提升多個任務(wù)的指標。我們主要考慮兩個在實際應(yīng)用中廣泛存在的約束條件：一，計算資源和存儲資源受限。此約束條件在某些硬件例如低功耗設(shè)備上格外突出。二，任務(wù)之間的耦合度應(yīng)當盡可能，使得每個任務(wù)可以在不影響其它任務(wù)的情況下得到迭代更新。

在傳統(tǒng)的遷移學(xué)習(xí)方案中，下游任務(wù)模型是在預(yù)訓(xùn)練模型的基礎(chǔ)上通過“微調(diào)”的方式獲得的[3]。在微調(diào)的過程中，任務(wù)模型以預(yù)訓(xùn)練模型的網(wǎng)絡(luò)權(quán)重作為初始權(quán)重，在任務(wù)標注數(shù)據(jù)上以梯度下降算法進行訓(xùn)練，逐步改變網(wǎng)絡(luò)權(quán)重，最終得到了一個新的模型。當有多個任務(wù)存在時，人們需要為每個任務(wù)微調(diào)得到一個獨一無二的任務(wù)模型，且這些任務(wù)模型中的權(quán)重和計算無法共用。顯而易見，傳統(tǒng)遷移學(xué)習(xí)方案計算復(fù)雜度和存儲復(fù)雜度都為O(n)，其中n 是任務(wù)的數(shù)量。這樣的復(fù)雜度在大多數(shù)應(yīng)用中是不可接受的。

2 模型設(shè)計

本文考慮擁有12 個Transformer 模塊的ViT-Base 模型。我們提出的方案流程由三個步驟構(gòu)成：（1）單任務(wù)微調(diào)訓(xùn)練；（2）單任務(wù)模型蒸餾；（3）多任務(wù)模型合并。模型方案的細節(jié)如下：

2.1 主干模型

2.1.1 Transformer 編碼器

本文使用的ViT-Base 模型擁有12 層Transformer 編碼器網(wǎng)絡(luò)。Transformer 模型在2017 年由A. Vaswani 等人提出，原用于機器翻譯任務(wù)。目前基于Transformer 結(jié)構(gòu)的模型已經(jīng)在許多自然語言處理和計算機視覺任務(wù)中取得了遠超傳統(tǒng)卷積神經(jīng)網(wǎng)絡(luò)（CNN）或長短時期記憶網(wǎng)絡(luò)（LSTM）的表現(xiàn)。Transformer 網(wǎng)絡(luò)有兩種類型，分別是編碼器和解碼器。一個Transformer 編碼器模塊由若干個編碼器和解碼器堆疊形成。如下圖的左側(cè)部分為編碼器，由多頭注意（Multi-Head Attention）和一個全連接組成，用于將輸入語料轉(zhuǎn)化成特征向量。右側(cè)部分是解碼器，其輸入為編碼器的輸出以及已經(jīng)預(yù)測的結(jié)果，保證訓(xùn)練與推理的一致性，用于輸出最后結(jié)果的條件概率，圖1。

圖1 Transformer 結(jié)構(gòu)圖

2.1.2 ViT 模型結(jié)構(gòu)

在計算機視覺有關(guān)的任務(wù)中，要讓Transformer 模型能夠接受圖像輸入，我們必須設(shè)法將一副二維圖像轉(zhuǎn)化成一個一維的向量序列。在ViT 模型中，我們把一副高度為H 寬度為W 的圖像拆分成N 個邊長為P 的正方形圖像小塊，其中N=HW/P2。若圖像中的每個像素包含C 個通道，則每個正方形小塊對應(yīng)一個維度為CP2 的向量，該向量記錄了圖像小塊的每個像素每個通道的原始數(shù)值。ViT 模型用一個可學(xué)習(xí)的全連接網(wǎng)絡(luò)將該向量非線性映射成了一個D 維的特征向量。通過這種方法，原始圖像變成了一個包含N 個D 維向量的序列，后者可以作為Transformer層的輸入。為了區(qū)分每個D 維向量對應(yīng)的正方形小塊在原始圖像中的位置，ViT 為每個位置的小塊引入了一個可學(xué)習(xí)的位置向量。該位置向量會加到特征向量中，作為后續(xù)Transformer 層的輸入。最后，ViT 引入了一個類似BERT 的[CLS]向量作為零號特征向量，該向量作為整幅圖像的特征表示。

2.1.3 ViT 模型預(yù)訓(xùn)練

我們使用的ViT-Base 模型在圖像網(wǎng)絡(luò)數(shù)據(jù)集上進行預(yù)訓(xùn)練。圖像網(wǎng)絡(luò)數(shù)據(jù)集是一個圖像分類數(shù)據(jù)集，其訓(xùn)練集、驗證集和測試集分別包含128 萬、5 萬和10 萬張標注圖片，分屬于1000 個類別。我們使用的ViT-Base模型在預(yù)訓(xùn)練任務(wù)上經(jīng)過300 次循環(huán)后最終可以達到84.15%的準確率。

2.2 多任務(wù)模型訓(xùn)練流程

2.2.1 單任務(wù)模型微調(diào)

第一個步驟，是針對每個任務(wù)微調(diào)預(yù)訓(xùn)練ViT 模型深層（靠近輸出層）的L 個模塊，凍結(jié)剩余位于淺層的12-L 個模塊。具體的L 數(shù)值是一個超參數(shù)，可以隨著任務(wù)的不同而變化。一個任務(wù)擁有較大的L 值意味著該任務(wù)被分配了更多的可訓(xùn)練參數(shù)以及計算資源。在極端情況下L=12，此時我們的微調(diào)方案和傳統(tǒng)的全微調(diào)等價。為了確保最低程度的參數(shù)和計算共享，我們在一個預(yù)先設(shè)定好的較小的區(qū)間N_min <= L <= N_max 內(nèi)搜尋最佳的L 值。在我們后續(xù)的實驗中，我們?nèi)_min=4 以及N_max=10。

第一個步驟最終為每一個任務(wù)輸出一個任務(wù)模型，這些任務(wù)模型的底層和預(yù)訓(xùn)練模型共享參數(shù)，頂層使用任務(wù)獨占的網(wǎng)絡(luò)參數(shù)。我們將這些模型稱為教師模型（見圖2 左），用于后續(xù)的模型蒸餾。

圖2 多任務(wù)ViT 模型設(shè)計圖

2.2.2 單任務(wù)模型蒸餾

在上一個步驟中，我們用微調(diào)部分網(wǎng)絡(luò)參數(shù)的辦法針對每個任務(wù)微調(diào)得到了一個任務(wù)模型。這些任務(wù)模型沒有直接的聯(lián)系，因此可以用知識蒸餾的方法壓縮這些模型的非共享網(wǎng)絡(luò)。具體的說，對一個第一步中得到的有著L_t 層任務(wù)獨占網(wǎng)絡(luò)以及12-L_t 層凍結(jié)網(wǎng)絡(luò)任務(wù)模型，我們可以用任意一種模型蒸餾算法將前者壓縮成一個僅有l(wèi)_t 層的更小的模塊，而讓后者保持不變(見圖2 中）。在實驗中，我們發(fā)現(xiàn)讓學(xué)生模型的網(wǎng)絡(luò)參數(shù)從教師模型的對應(yīng)網(wǎng)絡(luò)層中初始化可以達到最佳的學(xué)習(xí)效果。

2.2.3 多任務(wù)模型合并

在最后一步里，我們把所有的單任務(wù)學(xué)生模型合并成一個多任務(wù)模型，這樣凍結(jié)網(wǎng)絡(luò)層中的參數(shù)和所有計算可以被各個任務(wù)所共享，從而提升效率（見圖2 右）。在實際操作中，我們只需要在創(chuàng)建多任務(wù)模型計算圖后將各任務(wù)模型中的參數(shù)載入該計算圖中的對應(yīng)位置即可。

3 試驗方案

3.1 基線模型

在此我們將對比在一系列任務(wù)上我們的模型與主流基線模型的表現(xiàn)。

3.1.1 單任務(wù)模型（無模型蒸餾）：此模型由預(yù)訓(xùn)練模型ViT 在任務(wù)數(shù)據(jù)上以全參數(shù)微調(diào)的方式得到。此為預(yù)訓(xùn)練模型進行遷移學(xué)習(xí)的最佳實踐。如圖3 所示。

圖3 預(yù)訓(xùn)練ViT 模型結(jié)構(gòu)圖

3.1.2 單任務(wù)模型（有模型蒸餾）我們把擁有12 個Transformer 層的單任務(wù)模型通過模型蒸餾的方式壓縮至6 個模塊，從而降低計算復(fù)雜度。

3.1.3 多任務(wù)模型：采用傳統(tǒng)的多任務(wù)學(xué)習(xí)的方式訓(xùn)練，除了各任務(wù)輸出層以外共享全部的12 層傳遞參數(shù)，具體方案為每個任務(wù)分配任務(wù)專屬的1～3 個傳遞層用于學(xué)習(xí)任務(wù)相關(guān)的特征。所有任務(wù)共享預(yù)訓(xùn)練模型的前6 層網(wǎng)絡(luò)作為特征提取層。為了更加清晰的展示ViT 模型結(jié)構(gòu)和訓(xùn)練過程中的向量變化，圖4 給出了ViT的向量維度變化圖。

圖4 ViT 的向量維度變化圖

3.2 任務(wù)與數(shù)據(jù)集

3.2.1 分類任務(wù)：分類樣本包含100 個類別，每個類別共600 張圖片作為訓(xùn)練樣本，每張圖片的分辨率為32x32，如圖5 示例。

圖5 分類樣本示例

3.2.2 目標檢測：目標樣本包含20 萬張圖片，50 萬標注目標，共80 個類別。

3.2.3 語義分割：本樣本包含5000 張細粒度標注的圖片，2 萬張粗粒度標注的圖片，標注共分屬于30 個類別。

3.3 參數(shù)選擇與試驗

在接下來的實驗中使用Adam 的優(yōu)化器，采用學(xué)習(xí)率為2e-5 的32 批次訓(xùn)練超參數(shù)，從實驗結(jié)果中可以得出，我們的模型在圖像分類、語義分割和目標檢測這三個經(jīng)典的視覺任務(wù)上都超過了單任務(wù)模型蒸餾和多任務(wù)聯(lián)合訓(xùn)練的基線結(jié)果，僅僅只是略微遜于單任務(wù)模型全微調(diào)時的表現(xiàn)。僅使用3 層傳遞網(wǎng)絡(luò)，我們在分類任務(wù)、目標檢測和語義分割這三個數(shù)據(jù)集上任務(wù)上分別達到了12 層基線網(wǎng)絡(luò)指標的99.4%，98.3%，98.5%，如表1 所示。

表1 模型在三個視覺任務(wù)上的表現(xiàn)對比

4 模型分析

4.1 優(yōu)勢

本文提出的模型方案的主要優(yōu)點是它的靈活性。首先，不同的視覺任務(wù)可以從不同深度的網(wǎng)絡(luò)層獲取不同抽象層級的表征。其次，我們的方案可以允許給不同的任務(wù)分配不同的計算資源，例如更重要的任務(wù)可以分配更多的計算資源，或者降低資源投入邊際效費比較低的任務(wù)資源，從而使得模型在多任務(wù)環(huán)境下的平均表現(xiàn)打到更優(yōu)。第三，通過對任務(wù)之間的依賴進行解耦，我們的模型有效回避了傳統(tǒng)多任務(wù)模型的任務(wù)之間相互沖突的可能性，也不再需要針對任務(wù)權(quán)重超參數(shù)進行細致的調(diào)整。我們的多任務(wù)模型可以非常直觀的對單個任務(wù)子網(wǎng)絡(luò)模型的參數(shù)進行更新。

4.2 與適配器模型的比較

適配器模型是另一類在資源受限條件下應(yīng)用較為廣泛的模型。適配器模型通過在一個模型網(wǎng)絡(luò)的層與層之間添加一個小型可訓(xùn)練模塊（即適配器）來實現(xiàn)預(yù)訓(xùn)練模型對特定任務(wù)的適配。在遷移學(xué)習(xí)的過程中，原預(yù)訓(xùn)練模型中既有參數(shù)不發(fā)生變化，僅有適配器中的參數(shù)得到更新。這類方法的優(yōu)勢是只需要引入極少量的參數(shù)即可實現(xiàn)較優(yōu)秀的任務(wù)表現(xiàn)，但其劣勢在于計算開銷不會降低，因為任務(wù)網(wǎng)絡(luò)中的數(shù)據(jù)流被適配器完全改變了，無法跨任務(wù)復(fù)用。也就是說，在多任務(wù)場景下適配器模型的計算復(fù)雜度仍為O(n)，因此無法與本文中提出的方法進行直接的比較。

結(jié)束語

本文提出了一種新型多任務(wù)模型框架，可以在保證靈活性、低耦合性的同時極大降低模型對內(nèi)存的開銷和對計算資源的開銷。我們在多個視覺任務(wù)上證明了我們可以在保證模型預(yù)測精度的同時降低至多75%的資源開銷，因此具有較強的應(yīng)用推廣價值。