基于超大感受野注意力的超分辨率模型

注意力機(jī)制是深度學(xué)習(xí)領(lǐng)域非常重要的一個(gè)研究方向，在圖像超分領(lǐng)域也有不少典型的應(yīng)用案例，比如基于通道注意力構(gòu)建的RCAN，基于二階注意力機(jī)制構(gòu)建的SAN，基于像素注意力機(jī)制構(gòu)建的PAN，基于Transformer自注意力機(jī)制構(gòu)建的SwinIR，基于多尺度大核注意力的MAN等。

本文則以PAN為藍(lán)本，對(duì)其進(jìn)行逐步改進(jìn)以期達(dá)到更少的參數(shù)量、更高的超分性能。該方案具體包含以下幾個(gè)關(guān)鍵點(diǎn)：

• 提升注意力分割的感受野，類(lèi)似大核卷積注意力VAN；
• 將稠密卷積核替換為深度分離卷積，進(jìn)一步降低參數(shù)量；
• 引入像素規(guī)范化(Pixel Normalization)技術(shù)，其實(shí)就是Layer Normalization，但出發(fā)點(diǎn)不同。

上述關(guān)鍵技術(shù)點(diǎn)為注意力機(jī)制的設(shè)計(jì)提供了一個(gè)清晰的演變路線，最終得到了本文的VapSR，即大感受像素注意力網(wǎng)絡(luò)(VAst-receptive-field Pixel attention Network)。

實(shí)驗(yàn)結(jié)果表明：相比其他輕量超分網(wǎng)絡(luò)，VapSR具有更少的參數(shù)量。比如，項(xiàng)目IMDB與RFDN，VapSR僅需21.68%、28.18%的參數(shù)即可取得與之相當(dāng)?shù)男阅堋?/p>

本文動(dòng)機(jī)

通過(guò)引入像素注意力，PAN在大幅降低參數(shù)量的同時(shí)取得了非常優(yōu)秀的性能。相比通道注意力與空域注意力，像素注意力是一種更廣義的注意力形式，為進(jìn)一步的探索提供了一個(gè)非常好的基線。

受啟發(fā)于自注意力的發(fā)展，我們認(rèn)為：基于卷積操作的注意力仍有進(jìn)一步改進(jìn)的空間。因此，作者通過(guò)以下三個(gè)像素注意力中的設(shè)計(jì)原則展示了改善超分注意力的過(guò)程：

• 首先，在注意力分支引入大核卷積具有明顯的優(yōu)勢(shì)；
• 其次，深度分離卷積可以降低大核卷積導(dǎo)致的巨大計(jì)算復(fù)雜度問(wèn)題；
• 最后，引入像素規(guī)范化操作讓訓(xùn)練更高效、更穩(wěn)定。

**Large Kernel **以上圖i中的baseline為基礎(chǔ)，作者首先對(duì)注意力分支進(jìn)行感受野擴(kuò)增：將提升到(將圖示ii)，性能提升0.15dB，但參數(shù)量從846K提升到了4123K。

Parameter Reduction 為降低參數(shù)量，我們嘗試盡可能移除相對(duì)不重要的部分。作者提出了三個(gè)方案：(1) 將非注意力分支的卷積尺寸從下調(diào)到；(2) 將大核卷積注意力分支替換為深度深度分離卷積；(3) 將深度分離卷積中的深度卷積進(jìn)行分解為深度卷積+帶擴(kuò)張因子的深度卷積(該機(jī)制可參考下圖，將卷積拆分為+，其中后者的擴(kuò)張因子為3)。此時(shí)，模型性能變?yōu)?8.48dB，但參數(shù)量降到了240K，參數(shù)量基本被壓縮到了極限。

Pixel Normalization(PN) 注意力機(jī)制的元素乘操作會(huì)導(dǎo)致訓(xùn)練不穩(wěn)定問(wèn)題：小學(xué)習(xí)率收斂不夠好，大學(xué)習(xí)率又會(huì)出現(xiàn)梯度異常。前面的注意力改進(jìn)導(dǎo)致所得方案存在性能下降問(wèn)題。為解決該問(wèn)題，作者經(jīng)深入分析后提出了像素規(guī)范化技術(shù)(可參考下圖不同規(guī)范化技術(shù)的可視化對(duì)比)。

假設(shè)輸入特征為，第i個(gè)像素的特征均值與方差可以描述如下：

那么，像素規(guī)范化可以表示為：

當(dāng)引入PN后，模型的性能取得了顯著的提升，達(dá)到了28.92dB，參數(shù)量?jī)H為241K。

Switch Attention to Middle 在上述基礎(chǔ)上，作者進(jìn)一步將注意力的位置進(jìn)行了調(diào)整，放到了兩個(gè)卷積中間。此時(shí)，模型性能得到了0.03dB提升，達(dá)到了28.95dB，參數(shù)量仍為241K。

本文方案

前面的探索主要聚焦在微觀層面，基于此，作者進(jìn)一步在宏觀層面進(jìn)行了更多設(shè)計(jì)與提煉，進(jìn)而構(gòu)建了VapSR，取得了更佳的性能，同時(shí)具有更少的參數(shù)量。

上圖給出了所提VapSR架構(gòu)示意圖，延續(xù)了常規(guī)輕量方案的設(shè)計(jì)思路：

• 淺層特征：；
• 非線性映射： ;
• 圖像重建：

VAB模塊在前面探索得到的模塊上進(jìn)行了微調(diào)：(1) 主要是將模塊輸入與輸出通道數(shù)從64減少到了48，保持中間注意力部分的通道數(shù)仍為64；(2) 將注意力分支深度擴(kuò)張卷積(有時(shí)也稱之為空洞卷積)調(diào)整為深度擴(kuò)張卷積，此時(shí)感受野為；(3) 調(diào)整了注意力分支三個(gè)卷積的順序，將卷積移到最前面。對(duì)于VapSR-S，作者進(jìn)一步將部分從卷積調(diào)整為組卷積(group=2)，該操作可以進(jìn)一步降低參數(shù)量。

classAttention(nn.Module):
def__init__(self,dim):
super().__init__()
self.pointwise=nn.Conv2d(dim,dim,1)
self.depthwise=nn.Conv2d(dim,dim,5,padding=2,groups=dim)
self.depthwise_dilated=nn.Conv2d(dim,dim,5,1,padding=6,groups=dim,dilation=3)

defforward(self,x):
u=x.clone()
attn=self.pointwise(x)
attn=self.depthwise(attn)
attn=self.depthwise_dilated(attn)
returnu*attn

classVAB(nn.Module):
def__init__(self,d_model,d_atten):
super().__init__()
self.proj_1=nn.Conv2d(d_model,d_atten,1)
self.activation=nn.GELU()
self.atten_branch=Attention(d_atten)
self.proj_2=nn.Conv2d(d_atten,d_model,1)
self.pixel_norm=nn.LayerNorm(d_model)
default_init_weights([self.pixel_norm],0.1)

defforward(self,x):
shorcut=x.clone()
x=self.proj_1(x)
x=self.activation(x)
x=self.atten_branch(x)
x=self.proj_2(x)
x=x+shorcut

x=x.permute(0,2,3,1)#(B,H,W,C)
x=self.pixel_norm(x)
x=x.permute(0,3,1,2).contiguous()#(B,C,H,W)

returnx

本文實(shí)驗(yàn)

在實(shí)驗(yàn)部分，作者構(gòu)建了VapSR與VapSR-S兩個(gè)版本的輕量型超分方案：

• VapSR：包含21個(gè)VAB模塊，主干通道數(shù)為48；
• VapSR-S：包含11個(gè)VAB模塊，主干通道數(shù)為32。

此外，需要注意的是：對(duì)于X4模型，重建模塊并未采用常規(guī)的輕量方案(Conv+PS)，而是采用了類(lèi)EDSR的重方案(Conv+PS+Conv+PS)。

上表&圖給出了不同方案的性能與可視化效果對(duì)比，從中可以看到：

• 所提VapSR取得了SOTA性能，同時(shí)具有非常少的參數(shù)量。
• 在X4任務(wù)上，相比RFDN與IMDN，VapSR僅需21.68%/28.18%的參數(shù)量，即可取得平均0.187dB指標(biāo)提升；
• VapSR-S取得了與BSRN-S相當(dāng)?shù)男阅?，后者是NTIRE2022-ESR模型復(fù)雜度賽道冠軍。
• 在線條重建方面，VapSR具有比其他方案更精確的重建效果。