導(dǎo)讀

使用上下文注意力來(lái)進(jìn)行深度圖像修復(fù)。

今天，我們將深入探討深度圖像修復(fù)的一個(gè)突破，上下文注意力。通過(guò)使用上下文注意力，我們可以有效地從遙遠(yuǎn)的空間位置借用信息來(lái)重建局部缺失的像素。這個(gè)想法實(shí)際上或多或少和上一篇的復(fù)制-粘貼是一樣的。讓我們看看是如何做到的。

回顧

在之前的文章中，我介紹了使用shift連接層將未缺失區(qū)域生成的特征作為參考來(lái)恢復(fù)缺失區(qū)域，可以讓我們得到更好的修復(fù)結(jié)果。在這里，我們假設(shè)生成的特征是對(duì)ground truth的合理估計(jì)，并根據(jù)已知區(qū)域的特征與缺失區(qū)域內(nèi)生成的特征之間的相似性來(lái)確定合適的參考。

動(dòng)機(jī)

在圖像修復(fù)任務(wù)中，CNN的結(jié)構(gòu)不能有效地模擬缺失區(qū)域與遙遠(yuǎn)空間位置給出的信息之間的長(zhǎng)距離相關(guān)性。熟悉CNN的人應(yīng)該知道，在卷積層，核的大小和膨脹率控制著感受野，網(wǎng)絡(luò)需要越深入，才能看到整個(gè)輸入圖像。這意味著，如果我們想捕捉圖像的上下文，我們必須依賴(lài)于更深的層次，但我們丟失了空間信息，因?yàn)楦顚哟蔚奶卣鞯目臻g大小總是更小。因此，我們必須找到一種方法，在不用太加深網(wǎng)絡(luò)的情況下，從遙遠(yuǎn)的空間位置借用信息（即理解圖像的上下文）。

如果你還記得什么是膨脹卷積，你就會(huì)知道膨脹卷積是一種在早期的層中增加感受野而不添加額外參數(shù)的方法。然而，膨脹卷積有其局限性。它跳過(guò)連續(xù)的空間位置，以擴(kuò)大感受野。請(qǐng)注意，跳過(guò)的連續(xù)空間位置對(duì)于填充缺失的區(qū)域也很關(guān)鍵。

介紹

這項(xiàng)工作與我們以前討論過(guò)的網(wǎng)絡(luò)架構(gòu)、損失函數(shù)和相關(guān)技術(shù)類(lèi)似。對(duì)于該體系結(jié)構(gòu)，所提出的框架由兩個(gè)生成器網(wǎng)絡(luò)和兩個(gè)判別器網(wǎng)絡(luò)組成。這兩個(gè)生成器在全卷積網(wǎng)絡(luò)的基礎(chǔ)上使用了膨脹卷積。一個(gè)生成器用于粗重建，另一個(gè)用于細(xì)化。這被稱(chēng)為標(biāo)準(zhǔn)的從粗到細(xì)的網(wǎng)絡(luò)結(jié)構(gòu)。這兩個(gè)判別器同時(shí)在全局和局部看完整的圖像。全局判別器以整個(gè)圖像作為輸入，而局部判別器以填充區(qū)域作為輸入。對(duì)于損失函數(shù)，簡(jiǎn)單地說(shuō)，他們還使用了對(duì)抗損失（GAN損失）和L1損失（為了像素級(jí)重建精度）。對(duì)于L1損失，他們使用一個(gè)spatially discounted L1 loss，其中為每個(gè)像素差分配一個(gè)權(quán)值，權(quán)值基于像素到其最近的已知像素的距離。對(duì)于GAN損失，他們使用WGAN-GP損失，而不是我們所介紹的標(biāo)準(zhǔn)的對(duì)抗損失。他們聲稱(chēng)，這種WGAN對(duì)抗性損失也是基于L1距離度量，因此網(wǎng)絡(luò)更容易訓(xùn)練，訓(xùn)練過(guò)程也更穩(wěn)定。

在這篇文章中，我將專(zhuān)注于提出的上下文注意力機(jī)制。因此，我簡(jiǎn)要地介紹了從粗到細(xì)的網(wǎng)絡(luò)架構(gòu)、WGAN對(duì)抗損失和上面的加權(quán)L1損失。

方案

本文提出了上下文注意力機(jī)制，有效地從遙遠(yuǎn)的空間位置借用上下文信息來(lái)重建缺失的像素。將上下文注意力應(yīng)用到二次精細(xì)化網(wǎng)絡(luò)中。第一個(gè)粗重建網(wǎng)絡(luò)負(fù)責(zé)對(duì)缺失區(qū)域進(jìn)行粗估計(jì)。與前面一樣，使用全局和局部判別器來(lái)鼓勵(lì)生成的像素獲得更好的局部紋理細(xì)節(jié)。

貢獻(xiàn)

圖1，該模型在自然場(chǎng)景、人臉和紋理圖像上的修復(fù)效果實(shí)例。

本文最重要的思想是上下文注意力，它允許我們利用來(lái)自遙遠(yuǎn)空間位置的信息來(lái)重建局部缺失的像素。其次，使用對(duì)抗性損失和加權(quán)L1損失提高了訓(xùn)練的穩(wěn)定性。此外，本文提出的修復(fù)框架在自然場(chǎng)景、人臉、紋理等各種數(shù)據(jù)集上都獲得了高質(zhì)量的修復(fù)結(jié)果，如圖1所示。

方法

圖2，所提出修復(fù)方法的網(wǎng)絡(luò)架構(gòu)

圖2顯示了所提出修復(fù)框架的網(wǎng)絡(luò)架構(gòu)，如前所述，它由兩個(gè)生成器和兩個(gè)鑒別器組成。

上下文注意力

以下是這篇文章的主要關(guān)注點(diǎn)。讓我們來(lái)看看上下文注意力層是如何設(shè)計(jì)來(lái)借用遙遠(yuǎn)空間位置已知區(qū)域給出的特征信息來(lái)生成缺失區(qū)域內(nèi)的特征的。

圖3，上下文注意力層的圖形說(shuō)明圖3顯示了上下文注意層的圖形說(shuō)明。運(yùn)算是可微且全卷積的。

圖4，更多關(guān)于注意力注意力層的具體例子圖4是上下文注意層的一個(gè)更詳細(xì)的示例。如圖3所示，前景是指在缺失區(qū)域內(nèi)生成的特征，背景是指從已知區(qū)域中提取的特征。與復(fù)制粘貼方法類(lèi)似，我們首先要對(duì)缺失區(qū)域內(nèi)生成的特征與缺失區(qū)域外的特征進(jìn)行匹配。以圖4為例，生成的缺失區(qū)域內(nèi)的特征大小為64×64×64，假設(shè)缺失區(qū)域外的特征分為128個(gè)小特征patch，大小為64×3×3。注意，本例中特征的通道大小是64。然后，我們將128個(gè)小的feature patch與缺失區(qū)域內(nèi)生成的feature進(jìn)行卷積，得到大小為128×64×64的feature map。在本文中，該操作描述為：

其中{fx，y}為前景patches的特征， {bx‘y’}為背景patches的特征。sx，y，x‘，y’是缺失區(qū)域特征和已知區(qū)域特征之間的相似性。實(shí)際上，這是一個(gè)標(biāo)準(zhǔn)的余弦相似度計(jì)算過(guò)程。當(dāng)我們沿著通道維數(shù)看時(shí)，128個(gè)元素代表了所有已知patches和缺失區(qū)域內(nèi)特定位置之間的相似性。這反映了128個(gè)已知patches對(duì)該位置的貢獻(xiàn)。然后，我們沿著通道維度對(duì)特征映射執(zhí)行Softmax歸一化，如圖4中的藍(lán)色區(qū)域所示。在Softmax歸一化后，沿通道尺寸的每個(gè)位置之和應(yīng)為1。

與上一篇文章中提到的Shift-Net相比，你可以看到，這一次我們給每個(gè)已知特征的patch分配了權(quán)重，來(lái)表示重建的時(shí)候每個(gè)特征位置對(duì)于缺失區(qū)域的重要性（軟分配），而不是對(duì)于缺失區(qū)域的每個(gè)位置找一個(gè)最相似的（硬分配）。這也是為什么提出的上下文注意力是可微的。

最后，以注意力特征圖為輸入特征，以已知的patches為核，通過(guò)反卷積的方法重建缺失區(qū)域內(nèi)生成的特征。

注意力傳播

注意力傳播可以看作是注意特征圖的微調(diào)。這里的關(guān)鍵思想是，鄰近的像素通常有更接近的像素值。這意味著他們會(huì)考慮周?chē)h(huán)境的注意力值來(lái)調(diào)整每個(gè)注意力分?jǐn)?shù)。

例如，如果我們考慮左鄰居和右鄰居的注意力值，我們可以使用上面列出的公式更新當(dāng)前的注意力值。注意，k控制要考慮的鄰居的數(shù)量。作者聲稱(chēng)，這可以進(jìn)一步提高修復(fù)結(jié)果，這也可以通過(guò)與單位矩陣卷積作為核來(lái)實(shí)現(xiàn)。關(guān)于注意力機(jī)制的另一點(diǎn)是，采用了兩種技術(shù)來(lái)控制提取的已知特征塊的數(shù)量。

（i） 以較大的步長(zhǎng)提取已知的特征patch，以減少kernel數(shù)量。

（ii） 操作前先對(duì)特征圖大小進(jìn)行向下采樣，獲取注意力圖后再進(jìn)行上采樣。

網(wǎng)絡(luò)中的注意力

圖5，在第二個(gè)細(xì)化網(wǎng)絡(luò)中嵌入上下文注意力層的圖解

圖5顯示了作者如何將建議的上下文注意層整合到第二個(gè)細(xì)化網(wǎng)絡(luò)中。你可以看到，又引入了一個(gè)分支來(lái)應(yīng)用上下文注意力，然后將兩個(gè)分支連接起來(lái)以獲得最終的修復(fù)結(jié)果。注意力圖的顏色編碼是注意力圖的可視化方法。例如，白色意味著像素集中在自己身上，粉色是左下角區(qū)域，綠色是右上角區(qū)域，等等。你可以看到，這個(gè)例子有一個(gè)充滿(mǎn)粉紅色的注意力圖。這意味著填充區(qū)域從左下角區(qū)域借用了很多信息。

實(shí)驗(yàn)

作者首先比較了我們之前介紹的先前最先進(jìn)的技術(shù)。

圖6，比較提出的基線(xiàn)模型和GLCIC，從左到右，輸入圖像，GLCIC結(jié)果，基線(xiàn)結(jié)果圖6顯示了使用建議的基線(xiàn)模型和以前最先進(jìn)的GLCIC［2］進(jìn)行修復(fù)的結(jié)果。提出的基線(xiàn)模型如圖2所示，上下文注意力分支。很明顯，基線(xiàn)模型在局部紋理細(xì)節(jié)方面優(yōu)于GLCIC模型。請(qǐng)放大看清楚些。

圖7，對(duì)比基線(xiàn)和完整模型的修復(fù)結(jié)果。從左到右，ground truth，輸入圖像，基線(xiàn)結(jié)果，全模型結(jié)果，全模型注意圖圖7顯示了在Places2數(shù)據(jù)集上使用基線(xiàn)模型和完整模型（帶有上下文注意）的定性結(jié)果?？梢?jiàn)，完整模型具有較好的局部紋理細(xì)節(jié)，提供了較好的修復(fù)效果。這反映了上下文注意力層可以有效地從遙遠(yuǎn)的空間位置借用信息來(lái)幫助重建缺失的像素。請(qǐng)放大以便更好地觀(guān)看，特別是注意力圖。

表1，不同方法在Places2數(shù)據(jù)集上的定量比較表1列出了一些客觀(guān)的評(píng)價(jià)指標(biāo)，供參考。如前所述，這些指標(biāo)不能完全反映修復(fù)結(jié)果的質(zhì)量，因?yàn)橛性S多可能的解決方案來(lái)填補(bǔ)缺失的區(qū)域。你可以看到，建議的完整模型提供最好的L1， L2損耗和PSNR。對(duì)于電視丟失，PatchMatch提供更低的電視丟失，因?yàn)樗苯訌?fù)制原始圖像patch來(lái)填補(bǔ)漏洞。

供參考的全模型參數(shù)為2.9M。對(duì)于大小為512×512的圖像，GPU上每張圖像需要0.2秒，CPU上每張圖像需要1.5秒。

消融研究

注意力機(jī)制并不是一個(gè)新概念，文獻(xiàn)中有幾個(gè)注意力模塊。作者們用不同的注意力模塊做了實(shí)驗(yàn)。

圖8，通過(guò)使用不同的注意力模塊進(jìn)行修復(fù)。從左到右：輸入，使用空間變換網(wǎng)絡(luò)的結(jié)果，使用外觀(guān)流的結(jié)果，以及使用提出的上下文注意力的結(jié)果比較了文獻(xiàn)中兩個(gè)著名的注意力模塊，即空間變換網(wǎng)絡(luò)和外觀(guān)流。簡(jiǎn)單來(lái)說(shuō)，對(duì)于外觀(guān)流，使用卷積層代替上下文注意力層，直接預(yù)測(cè)二維像素偏移量作為注意力。這意味著我們添加一個(gè)卷積層來(lái)預(yù)測(cè)已知像素到缺失像素的移動(dòng)。在圖8中，你可以看到使用外觀(guān)流（中間）為不同的測(cè)試圖像提供類(lèi)似的注意力圖的結(jié)果。這就意味著注意力圖對(duì)于給予我們想要的“注意力”是沒(méi)有用的。你也可以觀(guān)察到空間變換網(wǎng)絡(luò)（左）不能為圖像修復(fù)任務(wù)提供有意義的注意力圖。一個(gè)可能的原因是空間變換網(wǎng)絡(luò)預(yù)測(cè)全局仿射變換的參數(shù)，這并不足以幫助填補(bǔ)缺失的區(qū)域，也需要局部信息。這里我沒(méi)有深入講解不同的注意力模塊。圖像修復(fù)中GAN損失的選擇。作者實(shí)驗(yàn)了不同的GAN損失，如WGAN損失，典型的對(duì)抗性損失，和最小平方GAN。他們通過(guò)經(jīng)驗(yàn)發(fā)現(xiàn)WGAN損失提供了最好的修復(fù)效果。重要的重建損失。在不使用L1損失的情況下訓(xùn)練了細(xì)化網(wǎng)絡(luò)。他們發(fā)現(xiàn)L1損失對(duì)于確保像素級(jí)重建精度是必要的，L1損失也會(huì)使修復(fù)結(jié)果變得模糊。因此，L1損失對(duì)于保證完整圖像更好的內(nèi)容結(jié)構(gòu)至關(guān)重要。

感知損失，風(fēng)格損失，TV損失。我們將很快說(shuō)到感知損失和風(fēng)格損失。一個(gè)簡(jiǎn)單的結(jié)論是，這三種損失并沒(méi)有給修復(fù)效果帶來(lái)明顯的改善。因此，他們的模型只使用加權(quán)L1損失和WGAN損失進(jìn)行訓(xùn)練。

總結(jié)

顯然，本文的核心思想是上下文注意力機(jī)制。上下文注意力層嵌入到第二個(gè)細(xì)化網(wǎng)絡(luò)中。注意，第一個(gè)粗重建網(wǎng)絡(luò)的作用是對(duì)缺失區(qū)域進(jìn)行粗略估計(jì)。這種估計(jì)用于上下文注意力層。通過(guò)匹配缺失區(qū)域內(nèi)生成的特征和缺失區(qū)域外生成的特征，我們可以知道缺失區(qū)域外所有特征對(duì)缺失區(qū)域內(nèi)每個(gè)位置的貢獻(xiàn)。注意，上下文注意力層是可微的和完全卷積的。

要點(diǎn)

你可能會(huì)發(fā)現(xiàn)，我們正越來(lái)越深入到深度圖像修復(fù)領(lǐng)域。我在上一篇文章中介紹了Shift連接層，它以硬分配的形式在CNN中嵌入了復(fù)制-粘貼的概念。本文以軟分配的形式構(gòu)造了一個(gè)上下文注意力層，該層是可微的，并且可以端到端學(xué)習(xí)，無(wú)需修改梯度的計(jì)算。希望大家能夠掌握本文提出的上下文注意力層的核心思想，特別是圖3和圖4所示的上下文注意力層公式。
編輯：lyn