利用ImageNet訓(xùn)練了一個(gè)能降噪、超分和去雨的圖像預(yù)訓(xùn)練模型

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

Motivation

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

準(zhǔn)備數(shù)據(jù)集

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

IPT

在上篇文章介紹過了，因?yàn)門ransformer本身是用于NLP領(lǐng)域的，輸入應(yīng)該是一個(gè)序列，因此這篇的論文做法和ViT[3]一樣，首先需要把feature map分塊，每個(gè)patch則視為一個(gè)word。但是不同的是，因?yàn)镮PT是同時(shí)訓(xùn)練多個(gè)task，因此模型定義了多個(gè)head和tail分別對應(yīng)不同的task。

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

Heads

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

Transformer encoder

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

fpi是feature map的一個(gè)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的輸入時(shí)encoder的輸出和task embedding。這些task embedding是可訓(xùn)練的，不同的task embedding代表處理不同的task。decoder的計(jì)算可以表示如下：

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

Tails

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

Loss

loss由兩部分組成，分別是Lcontrastive和Lsupervised的加權(quán)和。

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

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

實(shí)驗(yàn)與結(jié)果

作者用了32塊NVIDIA Tesla V100，以256的batch size訓(xùn)練了200個(gè)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.

責(zé)任編輯：lq