1. 程式人生 > 其它 >閱讀論文:《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》

閱讀論文:《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》

閱讀論文:《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》

來源:ICLR 2021 https://arxiv.org/abs/2010.11929
程式碼:https://github.com/google-research/vision_transformer

一、提出問題

這個工作保留了Transformer的基本結構和思想,遷移到了經典的影象分類問題。通過實驗,作者認為CNN並不是必須的,只用Transformer也能夠在分類任務中表現很好,尤其是在使用大規模訓練集的時候。同時,在大規模資料集上預訓練好的模型,在遷移到中等資料集或小資料集的分類任務上以後,也能取得比CNN更優的效能。

引入:

基於self-attention機制的Transformer結構是自然語言處理模型(NLP)的首選模型,其主要方法是在大型文字語料庫上進行預訓練,然後在較小的特定於任務的資料集上進行fine-tuning。

在計算機視覺中,卷積神經網路(ResNet)仍然占主導地位,此前的研究有受到NLP的啟發,嘗試將self-attention機制融合到CNN結構中,但由於使用了專門設定的注意力模式,因此尚未實現有效地擴充套件。

本文作者受到NLP中Transformer的啟發,通過將影象拆分為小patch,之後進行線性嵌入序列作為Transformer的輸入。最後以監督方式對模型進行影象分類任務的訓練。

二、主要思想

Transformer結構:

本文提出的模型:

Embadding:

使用了Transformer結構中的Encoder結構,由於標準的Transformer結構是一維的句子embadding,即二維矩陣[num_token, token_dim],所以需要對二維影象進行展平操作以適應結構。

以大小為224*224的影象為例,其資料格式為[H, W, C]三維格式,所以需要先通過Embedding層對資料做個變換。選取patch_size為16(ViT-B/16模型),輸入影象RGB三個維度的資料,對每一個維度按給定大小分成一堆Patches:(224/16)x(224/16)=196 個;其中每一個patch目前的shape為[16, 16, 3],通過一個線性對映將其展平為shape為[768]的向量

具體程式碼實現中使用了einops庫的Rearrange操作:

此外,在embadding後的影象序列中加了一個可學習的embedding向量專門用作分類,該序列在Transformer Encode的輸出可以用作得出影象分類的y。這個[class] token的格式和其他Token保持一致:shape為[768]的向量。之後和shape為[196,768]圖片Token序列進行Concat拼接操作得到[197, 768]維Token資料。

類同於Transformer中的Positional Encoding,此處也採用了一個可以訓練的Position Embedding,維度也是[197, 768]維。兩者進行add操作,最後得到的embadding為[197, 768]維。

Transformer Encoder:

由multi-head self-attention(MSA)和MLP塊組成,在每個塊之前應用Layer norm(LN),在每個塊之後應用殘差連線。MLP包含具有GELU和兩全連線層。

主要包括:

Norm層:此處使用了Layer Norm

殘差連線:參考Resnet的結構

多頭注意力機制:Multi-Head Attention,類同於Tramsformer

MLP Block:全連線+GELU啟用函式+Dropout

最後輸出的資料維度仍然是[197,768]維。

MLP Head:

對於分類任務,只需要提取出[class]token生成的對應結果即可,即[197, 768]維中抽取出[class]token對應的[1, 768]。MLP Head主要就是一個普通的分類結構:Layer Norm後加一個Linear,最後softmax分類即可。

訓練採取先在大資料集上預訓練,再遷移到小資料集上面。做法是把ViT的MLP Head去掉,換成一個768 x K的全連線層。其中K為對應資料集的類別數。

三、實驗

實驗一:效果測試

效能超過了CNN的SOTA結果(對比:Big Transfer (BiT)模型和Noisy Student模型)

實驗二:對預訓練資料的要求

資料集擴大後,ViT模型的優勢顯著。

實驗三:注意力機制

四、存在的問題以及改進

CNN模型在影象任務中會有一些優點,比如平移不變性和區域性性等,這些特點Transformer並不具備,因此該方法在訓練不足的資料量時不能很好地學習,從結果上來看對中等規模的資料集進行訓練時,精度要低於同等大小的ResNet系列模型。但對於更大規模的資料集(JFT-300)上的訓練,CNN模型可能會趨於飽和,Transformer的效能要勝過inductive bias。但是首先JFT-300是私有的資料集,很難復現結果;此外,Transformer結構若想取得理想的效能和泛化能力,計算資源和資料集都不具備,很難有所收穫。

對此的解決:《Training data-efficient image transformers & distillation through attention》

使用一個性能很好的分類器 (CNN/Transformer,實驗驗證) 作為teacher model,設定了distillation token和蒸餾損失,通過知識蒸餾的方法進行訓練。

主要改進:引入了一個distillation token,然後在自注意力機制中跟class token、patch token不斷互動學習知識,這個distillation token要跟teacher model預測的label一致。

五、擴充套件

基於這個思想,傳統的CNN網路也可以和Transformer結合起來,將其他網路尾部的特徵取出來進行Patch embadding,然後按照類似的方法進行組合和訓練。

Vision Transformer在其提出的一年期間,很多應用於影象檢索、ReID、自監督學習等領域的Vision Transformer快速被提出,並取得了較為優異的成績。

多目標跟蹤 (MOT):《TrackFormer: Multi-Object Tracking with Transformers》

3D:《3D-Transformer: Molecular Representation with Transformer in 3D Space》

YOLOS:《You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detection》

點雲:《PoinTr: Diverse Point Cloud Completion with Geometry-Aware Transformers》

Neural Architecture Search:《GLiT: Neural Architecture Search for Global and Local Image Transformer》

自動駕駛:《Multi-Modal Fusion Transformer for End-to-End Autonomous Driving》

文字到視覺檢索:《Thinking Fast and Slow: Efficient Text-to-Visual Retrieval with Transformers》

基於自然語言查詢視訊幀片段:《QVHighlights: Detecting Moments and Highlights in Videos via Natural Language Queries》

掩碼自編碼器:《Masked Autoencoders Are Scalable Vision Learners》