圖1 本文方法和現(xiàn)有方法的一些對(duì)比結(jié)果

本文介紹我們?cè)贑VPR 2022發(fā)表的用于圖像修復(fù)的工作。該工作是基于當(dāng)前流行的Transformer實(shí)現(xiàn)的，目的是減少Transformer在應(yīng)用到圖像修復(fù)過程中的一些信息損失，從而提升模型修復(fù)圖片的質(zhì)量。目前論文和代碼都已經(jīng)公開，歡迎大家試用交流。

論文標(biāo)題: Reduce Information Loss in Transformers for Pluralistic Image Inpainting

作者單位: 中國科學(xué)技術(shù)大學(xué)，微軟云AI

錄用信息: CVPR 2022

代碼鏈接: https://github.com/liuqk3/PUT

論文鏈接: https://arxiv.org/pdf/2205.05076

一、針對(duì)問題

本文主要的任務(wù)是圖像修復(fù)，即輸入的圖片是殘缺的，通過算法根據(jù)提供的殘缺圖片來補(bǔ)全圖片中殘缺的部分。早期的一些工作所利用的模型，一般都是基于CNN實(shí)現(xiàn)的。由于CNN具有一定的位置偏置等因素，導(dǎo)致其修復(fù)的圖片可視化效果并不理想。近兩年，隨著Transformer在計(jì)算機(jī)視覺領(lǐng)域的巨大成功，一些研究人員開始利用Transformer來進(jìn)行圖像修復(fù)，并取得了非常不錯(cuò)的效果。然而，通過對(duì)這些方法進(jìn)行分析，發(fā)現(xiàn)這類方法普遍存在或多或少的信息損失，原因主要有兩點(diǎn)：（1）對(duì)圖片進(jìn)行下采樣。眾所周知，Transformer的計(jì)算量與輸入的序列長度呈平方關(guān)系。為了減少計(jì)算量，要保證輸入Transformer的序列長度在可接受范圍內(nèi)。為此，現(xiàn)有方法會(huì)將圖片進(jìn)行下采樣，比如從256x256下采樣到32x32。這種下采樣的方式不可避免地帶來了信息損失；（2）量化。像素的個(gè)數(shù)是256^3，如果直接將每個(gè)像素（實(shí)際上是索引）都當(dāng)作一個(gè)token，那么Transformer內(nèi)部的embedding個(gè)數(shù)也是256^3，較多的embedding不僅帶來了較多的參數(shù)，也不利于模型訓(xùn)練。因此現(xiàn)有方法會(huì)對(duì)像素進(jìn)行聚類量化，比如從256^3個(gè)變成512個(gè)。量化的過程也會(huì)帶來信息損失。

本文的方法旨在解決上述的問題。（1）利用auto-encoder代替下采樣。auto-encoder中的encoder會(huì)將輸入的圖片分成獨(dú)立的圖片塊，每塊獨(dú)立編碼，避免信息交叉影響。Decoder負(fù)責(zé)從輸入的量化的特征重構(gòu)圖片；（2）不量化。為了避免Transformer中的信息損失，輸入到Transformer中的不是離散的token（即索引），而是encoder出來的特征。離散的token只用作Transformer的輸出。圖2對(duì)比了本文方法和現(xiàn)有方法的流程。

二、方法框架

本文方法的具體流程如圖3所示。整個(gè)算法包含兩部分：auto-encoder和Transformer。其中auto-encoder是經(jīng)過特殊設(shè)計(jì)，專門用于圖像修復(fù)任務(wù)。Transformer也與現(xiàn)有的方法中的Transformer（如DALL-E, IGPT等）稍有不同。下面將分別進(jìn)行介紹。

1. Auto-encoder

本文方法中的auto-encoder全稱為Patch-based Vector Quantized Variational Auto-Encoder(P-VQVAE)。其encoder全部由全連接層構(gòu)成。輸入的殘缺圖片會(huì)被分成獨(dú)立的圖片塊，每一塊都單獨(dú)經(jīng)過encoder進(jìn)行獨(dú)立編碼。不同的圖片塊缺失的像素個(gè)數(shù)也不同。獨(dú)立編碼避免了相互之間的干擾，利用提取更具有區(qū)分度的特征。

Encoder輸出的特征經(jīng)過域碼本量化之后輸入到decoder可以實(shí)現(xiàn)圖像重構(gòu)。Decoder是由conv層組成，其不同之處在于有一個(gè)參考分支。參考分支的輸入就是提供的殘缺圖片，目的是保證殘缺圖片中已有的像素值保持不變。當(dāng)參考分支不使用時(shí)，P-VQVAE可以實(shí)現(xiàn)和現(xiàn)有auto-encoder一樣的作用，即圖像重構(gòu)。

P-VQVAE的訓(xùn)練和VQVAE基本一致，不過我們引入了其他的損失來提高其重構(gòu)的性能，如引入判別器對(duì)抗訓(xùn)練等。

2. Transformer

本方法中的Transformer全稱是Un-Quantized Transformer。它的輸入是encoder輸出的特征，而不是離散的token。這種做法的目的是避免量化引入信息損失。對(duì)于每個(gè)圖片塊，Transformer會(huì)輸出一個(gè)概率，這個(gè)概率就是圖片塊對(duì)應(yīng)碼本中的特征的概率。在訓(xùn)練階段，只需要一個(gè)簡單測(cè)交叉熵?fù)p失即可。在測(cè)試階段，只要圖片塊中的任意一個(gè)像素缺失，就會(huì)其對(duì)應(yīng)概率進(jìn)行采樣，將從碼本中采樣得到特征向量作為該圖片塊的特征，并送入decoder用于圖像重構(gòu)。

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

與相關(guān)工作的定量對(duì)比如圖4所示，可以看到本文的方法能夠取得非常不錯(cuò)的結(jié)果，尤其是在ImageNet數(shù)據(jù)集上的效果更好。然而對(duì)于圖像修復(fù)/生成類的任務(wù)，定量的指標(biāo)很難反應(yīng)算法的優(yōu)劣，為此我們提供了大量的可視化結(jié)果，分別如圖5，6，7，8所示。

四、總結(jié)

本文方法針對(duì)圖像修復(fù)任務(wù)設(shè)計(jì)了一個(gè)全新的框架，主要是減少現(xiàn)有方法在利用Transformer進(jìn)行圖像修復(fù)時(shí)的信息損失問題。通過實(shí)驗(yàn)可以看到，本文方法在指標(biāo)、可視化效果上的提升是非常顯著的。但是基于Transformer實(shí)現(xiàn)的自回歸模型，在測(cè)試階段時(shí)的速度都比較慢，本文設(shè)計(jì)的PUT也有這個(gè)問題。但是在測(cè)試階段采用自回歸的形式，主要目的是得到多樣化的結(jié)果。當(dāng)只需要得到一個(gè)結(jié)果時(shí)，實(shí)際上可以通過一次網(wǎng)絡(luò)前傳，得到所有的token，進(jìn)而大大提高測(cè)試的速度。另外，采樣策略對(duì)修復(fù)圖片的結(jié)果影響還是很大的，因此如何設(shè)計(jì)更加有效的采樣策略，是一個(gè)值得進(jìn)一步討論的問題。

公眾號(hào)后臺(tái)回復(fù)“CVPR 2022”獲取論文合集打包下載～