一種基于遷移學習的視覺多任務模型探析

劉夏鳴

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

1 概述

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

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

2 模型設計

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

2.1 主干模型

2.1.1 Transformer 編碼器

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

圖1 Transformer 結構圖

2.1.2 ViT 模型結構

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

2.1.3 ViT 模型預訓練

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

2.2 多任務模型訓練流程

2.2.1 單任務模型微調

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

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

圖2 多任務ViT 模型設計圖

2.2.2 單任務模型蒸餾

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

2.2.3 多任務模型合并

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

3 試驗方案

3.1 基線模型

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

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

圖3 預訓練ViT 模型結構圖

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

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

圖4 ViT 的向量維度變化圖

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

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

圖5 分類樣本示例

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

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

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

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

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

4 模型分析

4.1 優(yōu)勢

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

4.2 與適配器模型的比較

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

結束語

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