利用ImageNet訓練了一個能降噪、超分和去雨的圖像預訓練模型

說到Transformer，大家可能會想到BERT[1]、GPT-3[2]等等，這些都是利用無監(jiān)督訓練的大型預訓練模型。既然Transformer也能用在CV上，那么能不能做類似的事情呢？這篇論文利用ImageNet訓練了一個能降噪、超分和去雨的圖像預訓練模型（IPT）。

Motivation

目前很多l(xiāng)ow-level的task其實都是有一定相關性的，就是在一個low-level task上預訓練對另一個task是有幫助的，但是目前幾乎沒有人去做相關的工作。而且pre-training在某些數(shù)據稀缺的task上就很有必要，并且無論在CV還是NLP，使用pre-trained model是非常常見的事情。對于一些輸入和輸出都是image的low-level算法來說，目前的pre-trained model顯然是不適合的。

準備數(shù)據集

因為Transformer需要大量的數(shù)據去擬合，所以必須使用一個大型的數(shù)據集。在這篇論文中，作者用的是imagenet。對于imagenet的每一張圖片生成各種任務對應的圖像對，例如對于超分（super-resolution）來說，模型的輸入數(shù)據是imagenet經過下采樣的數(shù)據，而標簽是原圖。

IPT

在上篇文章介紹過了，因為Transformer本身是用于NLP領域的，輸入應該是一個序列，因此這篇的論文做法和ViT[3]一樣，首先需要把feature map分塊，每個patch則視為一個word。但是不同的是，因為IPT是同時訓練多個task，因此模型定義了多個head和tail分別對應不同的task。

整個模型架構包含四個部分：用于提取特征的heads、Transformer Encoder、Transformer Decoder和把feature map還原成輸出的tails。

Heads

不同的head對應于不同的task，由于IPT需要處理多個task，因此是一個multi-head的結構，每個head由3層卷積層組成。Heads要完成的任務可以描述為：fH = Hi(x)，x是輸入圖像，f是第i個Head的輸出。

Transformer encoder

在輸入Transformer前，需要將Head輸出的feature map分成一個個patch，同樣還需要加入位置編碼信息，與ViT不同，這里是直接相加就可以作為Transformer Encoder的輸入了，不需要做linear projection。

fpi是feature map的一個patch，Epi∈ RP*P×C是fpi的learnable position encoding。LN是layer normalization，MSA是多頭self-attention模塊，F(xiàn)FN是feed forward network。

Transformer decoder

Transformer decoder的輸入時encoder的輸出和task embedding。這些task embedding是可訓練的，不同的task embedding代表處理不同的task。decoder的計算可以表示如下：

fEi是指encoder的輸出，fDi是指decoder的輸出。

Tails

Tails與Heads是相對應的，但是不同的tail的輸出的大小可能不一樣，例如超分，做超分時輸出比輸入的圖像大，因此與其它的tail輸出的大小可能不一樣。

Loss

loss由兩部分組成，分別是Lcontrastive和Lsupervised的加權和。

Lsupervised是指IPT的輸出與label的L1 loss。

加入Lcontrastive是為了最小化Transformer decoder對于來自同一張圖的不同patch的輸出的距離，最大化對于不同圖片的patch之間的輸出的距離。

實驗與結果

作者用了32塊NVIDIA Tesla V100，以256的batch size訓練了200個epoch。

Reference

[1]Jacob Devlin, Ming-Wei Chang, Kenton Lee, and KristinaToutanova. Bert: Pre-training of deep bidirectionaltransformers for language understanding. arXiv preprintarXiv:1810.04805, 2018.

[2]Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.Language models are few-shot learners. arXiv preprintarXiv:2005.14165, 2020.

[3]Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale[J]. arXiv preprint arXiv:2010.11929, 2020.

責任編輯：lq