基于局部期望最大化注意力的圖像降噪

李澤田，雷志春

(天津大學 微電子學院，天津 300072)

1 引 言

圖像降噪是計算機視覺中的基本問題，也是許多其他圖像恢復任務的重要預處理步驟[1-6]。通常，圖像降噪的目的是從其退化圖像恢復出無噪聲圖像[7]。

圖像降噪具有非適定性(Ill-posed)[8]。傳統圖像降噪方法[9-12]應用各種約束來模擬噪聲特征(例如，均勻、非均勻、深度感知)，并利用不同的圖像先驗來規范解空間。然而，在大多數方法中使用的圖像先驗都是手工制作的,并且消耗大量的計算資源。

最近幾年，隨著深度卷積神經網絡(Deep Convolutional Neural Network，DCNN)的發展，圖像去噪性能得到了顯著提升。文獻[13]通過將批量歸一化集成到殘差網絡中，提出了一種稱為DnCNN(Denoising Convolutional Neural Networks)的有效圖像降噪網絡，該網絡優于傳統的降噪算法。文獻[14]基于小波分解，提出了一種基于多級小波的降噪算法，稱為MWCNN(Multi-level Wavelet-CNN)。文獻[15]提出了一種快速靈活的網絡稱為FFDNet，可以處理具有不均勻噪聲的圖像。為了利用DCNN中圖像特征的非局部性，文獻[16]在去噪網絡中采用k-最近鄰匹配，提出N3Net算法。

由于DCNN強大的擬合能力，上述基于DCNN的降噪算法較傳統算法效果提升顯著。但這些算法也有其局限性。首先，采用“編碼器-解碼器”級聯的方式構建模型，在計算每一個解碼器的狀態時需要關注所有的編碼器輸入，計算量比較大，靈活性差，網絡的泛化能力較低。其次，平等地處理不同通道、不同區域的特征，效率較低，浪費有限的計算資源。最后，在DCNN中，隨著圖像在網絡中的傳播，特征域的噪聲會不斷累加[17]。

為克服以上基于DCNN的降噪算法的缺點，本文設計了一個基于局部期望最大化注意力的圖像降噪算法。近年來，注意力機制在自然語言處理領域取得卓越成果[18]。Nonlocal[19]第一次將注意力機制應用到計算機視覺領域。RCAN[20]提出后，注意力在低層視覺領域也受到了廣泛的關注。注意力機制關注圖像中突出的特征(紋理細節等)，并使每個像素充分利用不同通道的語義信息，有利于提升網絡效率。然而，這些注意力機制需要生成一個巨大的注意力圖，其空間復雜度和時間復雜度巨大。其瓶頸在于，每一個像素的注意力圖都需要對全圖計算。期望最大化注意力機制(Expectation-Maximization Attention，EMA)[21]，摒棄了在全特征圖上計算注意力圖的流程，首先通過期望最大化(EM)算法迭代出一組緊湊的基，然后在這組基上運行注意力機制，大幅降低了計算復雜度。然而，EMA仍然是針對全圖計算期望值，這對于降噪任務不是最優方案。因為噪聲屬于快速變化的高頻信號，本文認為，應該利用特征圖的局部信息提取噪聲注意力圖，即采用局部期望最大化(Local Expectation-Maximization Attention，LEMA)算法替代EMA算法，通過局部期望最大化(LEM)算法迭代出一組緊湊的局部基，在這組局部基上運行注意力機制，這樣不僅可以有效濾除噪聲，還可以進一步降低復雜度。本文把這一機制嵌入網絡中，構造出輕量且易實現的 LEMA Unit，在多個數據集上取得了較好的降噪效果。本文的主要貢獻如下：

(1)改進了期望最大化算法，提出了更適于降噪任務的局部期望最大化算法(LEMA)，該算法對輸入的局部特征進行迭代得到一組緊湊的局部基，在這組局部基上計算注意力圖可以大幅降低計算復雜度。

(2)將LEMA機制引入DCNN中，結合通道注意力與空間注意力，構造出LEMA Unit，并將LEMA Unit嵌入到自動編碼器網絡中，提出LEMANet模型。

本文進行了大量的圖像降噪實驗，以驗證LEMANet的有效性。 結果表明，LEMANet在視覺質量和定量測量方面均實現了領先的降噪性能。

2 基于DCNN的圖像降噪

近幾年的時間里，卷積神經網絡發展的一個趨勢是為了獲得較好的性能，不斷地增加網絡層數。DCNN通過堆疊多個簡單的非線性單元從大量標記數據中學習高度非線性映射[22]。

自動編碼器(Autoencoder)是一種無監督的學習算法，主要用于數據的降維或者特征的提取，由編碼器(Encoder)和解碼器(Decoder)兩部分組成，每一部分均采用堆疊DCNN的方式構成。編碼器將輸入圖像(設為x)經過加權W、映射(φ)之后提取到圖像主特征y：

(1)

其中：L為編碼器層數。

解碼器對y反向加權WT、映射(φ)得到輸出z，

(2)

其中：L為解碼器層數。

反向傳播過程通過反復迭代訓練兩組權值W、WT，使得損失函數(MSE，感知損失等)最小，盡可能保證z近似于x，即可重構輸入圖像x。自動編碼器結構如圖1所示。fθ是編碼器映射，由于輸入數據x中存在相關性，為了提取圖像表征的關鍵特征，在編碼過程施加“瓶頸”(Bottleneck)，即對輸入數據進行壓縮，提取得到關鍵特征y。解碼器負責由特征y解碼恢復出重構數據z。在文獻[23-25]中，采用降噪自動編碼器(Denoising Autoencoder，DAE)方法進行高斯噪聲去除，并獲得與BM3D[26]、K-SVD[27]等傳統算法相當的結果，顯示了自動編碼器方法的降噪潛力。

圖1 自動編碼器結構圖Fig.1 Structure of autoencoder

3 局部期望最大化注意力機制

3.1 局部期望最大化算法

期望最大化(Expectation-Maximization，EM)[28]算法旨在找到隱含變量模型的最大似然解。用Y表示觀測變量的數據，Z表示隱含變量的數據，Y和Z一起構成完全數據。假設觀測數據Y概率分布為P(Y|θ)，其中θ是模型待估計的參數，Y和Z的聯合概率分布為P(Y,Z|θ)。一般情況下，Z的分布未知，在θ給定的情況下，通過后驗概率P(Z|Y,θ)迭代得到。完全數據的似然函數為lnP(Y,Z|θ)，EM算法通過迭代的方法求lnP(Y,Z|θ)的極大似然估計，每次迭代分兩步：E步和M步，不斷重復E步和M步，直到收斂。

在E步，根據給定的θ通過后驗概率P(Z|Y,θ)求得Z的分布，然后求完全數據的似然函數的數學期望EZ|Y,θ：

(3)

式中：θi為第i次迭代的模型參數，θi+1為第i+1次迭代的模型參數。

在M步，求使式(3)極大化的θi，確定第i+1迭代的參數估計值θi+1，

θi+1=argmaxθEZ|Y,θi，

(4)

式中argmax為極大值函數。

考慮到圖像不同區域的噪聲分布并不一定是均勻的，所以本文提出在圖像局部應用期望最大化機制，其結構如圖2所示。首先采用邊緣檢測，區分不同的噪聲區域[29]。然后在每一區域內部應用期望最大化機制，迭代出一組緊湊的局部基，為后續計算注意力特征圖做準備。最后，在局部基上計算注意力特征圖：

LEM(z)=φlocal([E(z),M(z)])，

(5)

其中：LEM(z)為局部基，φlocal為局部噪聲檢測算子，[E(z),M(z)]為期望最大化算子。

圖2 局部期望最大化注意力模塊Fig.2 Local expectation-maximization attention module

3.2 通道注意力

網絡的淺層特征包含豐富的高頻信息，這對降噪后圖像細節的恢復非常重要，但這些高頻信息在特征通道間被平等對待，這阻礙了網絡的表達能力。為克服這一缺點，本文將通道注意力機制引入到圖像降噪中。

通道注意力主要關注輸入的抽象信息[30]。輸入圖像各特征通道間的信息有高度相關性，利用這一相關性，濾除通道冗余信息，使網絡專注于關鍵信息。生成通道注意力，不僅可以提高網絡性能，還可以降低網絡復雜度。為了有效計算通道注意力，需要對輸入特征圖進行降維，即減少冗余信息，主要是低頻信息。常用的降維方法是利用均值池化[31-32]，但是均值池化在濾除低頻的同時，會損失一些高頻信息。為此，本文同時使用了均值池化與最大池化，實驗證明效果比單獨使用均值池化性能優異。通道注意力滿足方程：

(6)

圖3 視覺質量對比。圖像噪聲水平為25的降噪結果。Fig.3 Visual quality comparison. Denoising results of one image from BSD68 with noise level 25.

圖4 視覺質量對比。圖像噪聲水平為50的降噪結果。Fig.4 Visual quality comparison. Denoising results of one image from BSD68 with noise level 50.

圖5 通道注意力模塊結構Fig.5 Structure of channel attention moduel

3.3 空間注意力

MS(F)=σ(f7×7
([AvgPool(F);MaxPool(F)]))，

(7)

式中：σ為Sigmoid函數；f7×7() 為7×7的卷積操作；AvgPool()為均值池化函數；MaxPool()為最大池化函數；[;]表示通道維度上特征的連接操作。

空間注意力模塊的結構如圖6所示。

圖6 空間注意力模塊結構Fig.6 Structure of spatial attention moduel

3.4 全局注意力網絡

在卷積神經網絡中，設卷積核為n×n，則該層卷積只關注鄰域為n的區域，即使通過增加層數、采用擴張卷積使感受野越來越大，也難以學習全局語義信息[19]。全局注意力網絡就是提取全局特征，也就是圖像中任何像素在當前像素的映射的疊加。

全局注意力滿足方程：

(8)

全局注意力網絡如圖7所示，其結構由3部分組成：(1)注意力生成，采用1×1卷積和softmax函數來獲取注意力權重，然后與輸入特征圖相乘得到全局注意力特征；(2)特征傳遞，此部分采用1×1卷積、實例歸一化和非線性激活層(ReLU)級聯的方式；(3)全局注意力聚合，得到全局注意力。該模塊中的1×1卷積可以實現特征數的壓縮，從而降低計算復雜度。

圖7 全局注意力模塊結構Fig.7 Structure of global attention network

4 局部期望最大化注意力網絡

局部期望最大化注意力網絡，利用局部期望最大化機制迭代出一組緊湊的局部基，然后在這組局部基上利用通道注意力模塊提取圖像的抽象信息，即圖像中的內容；利用空間注意力模塊提取圖像中的位置信息，即圖像中各要素間的相對位置。全局注意力模塊負責計算圖像中所有像素點對于某一特定像素的映射，其中全局池化操作可以提高網絡的擬合能力，減少網絡層數，加速訓練。

局部期望最大化注意力網絡由5部分組成：(1)淺層特征提取模塊，由一個3×3的卷積層構成，用于提取淺層特征，這些特征包含大量的高頻信息；(2)深層特征提取模塊，用于提取深層特征，這些特征包含抽象信息，由一個自動編碼器Autoencoder構成，Autoencoder采用的是Encoder-Decoder結構；(3)局部期望最大化模塊，用于迭代出一組緊湊的局部基；(4)注意力模塊，由通道注意力、空間注意力網絡及全局注意力網絡構成；(5)圖像重建模塊，由一層卷積核大小為1×1的卷積層構成，因為生成的是彩色圖像，所以輸出通道為3。整個網絡結構如圖7所示。

圖8 全局雙注意力網絡結構Fig.8 Structure of global dual attention network

通道注意力模塊與空間注意力模塊中采用最大池化與均值池化操作,突出主要特征，抑制次要特征，生成通道注意力與空間注意力。全局注意力模塊采用全局池化提取圖像全局特征，生成全局注意力。由于抑制了對降噪任務貢獻較小的次要特征，因而使用注意力機制降低了特征圖計算復雜度。

5 測量實驗與結果

為了驗證本文推出的局部期望最大化注意力網絡的降噪性能，進行了多次實驗，并與其他現有的降噪模型進行了全面對比，包括深度學習和非深度學習的方法：BM3D[26]、WNNM[12]、DnCNN[13]、FFDNet[15]、N3Net[16]。

5.1 實驗參數與數據集

5.1.1 參數設置

淺層與深層特征提取模塊，卷積核尺寸為3×3，歩長為1，特征通道數為64；全局特征提取模塊，卷積核尺寸為1×1，歩長為1，特征通道數為128；圖像重建模塊，卷積核尺寸為1×1，歩長為1，特征通道數為3。

訓練集原圖使用DIV2K數據集[33]，該數據集來自CVPR NTIRE 2017挑戰的高質量(2K分辨率)圖像數據集。DIV2K數據集包含800個訓練圖像，100個驗證圖像和100個測試圖像。噪聲圖則依據文獻[14]，生成3種噪聲等級的噪聲圖，即σ=15，25，50。測試集使用3個標準數據集：Set12[17]、BSD68[34]和Urban100[35]。實驗結果測試指標在YCbCr空間的Y通道上用峰值信噪比(PSNR)和結構相似性指數(SSIM)[36]評價。

5.1.2 實驗細節

對800個訓練圖像進行數據增強，訓練圖像隨機旋轉90°，180°，270°并進行水平翻轉。在每個訓練批次中，提取16對64×64的原圖與噪聲圖作為輸入。訓練過程中，使用ADAM優化器[37]最小化損失函數。初始學習率設置為1E-4，并隨訓練過程呈指數衰減，衰減系數為0.9。

5.2 實驗參數與數據集

與5個最先進算法進行了對比實驗：BM3D[26]、WNNM[12]、DnCNN[13]、FFDNet[15]、N3Net[16]。測試集為Set12[13]、BSD68[34]和Urban100[35]。實驗結果如表1所示。可以看出，本文提出的LEMA的PSNR值比其他最先進算法均有顯著提升。σ=25時，對于測試集Set12，LEMA的PSNR值比傳統方法BM3D和WNNM提升0.36～0.67 dB左右，比DnCN提升0.20 dB左右，比FFDNet提升0.21 dB，比N3Net提升0.06 dB。在數據集BSD68與Urban100上，LEMA均表現出較明顯的提升。

視覺質量對比實驗如圖3、圖4所示。可以看出，相較于其他算法，LEMA的降噪效果明顯，保留更多的細節信息。

表1 與最先進算法對比實驗Tab.3 Performance comparison with the state of the art

6 結 論

本文提出了一種基于局部期望最大化注意力的圖像降噪網絡，該網絡利用期望最大化在圖像局部迭代出局部基，然后在局部基上采用注意力網絡，在保證精度的同時，大大降低特征圖的計算復雜度。與最先進的降噪算法對比顯示該算法的性能提升明顯。σ=25時，對于測試集Set12，LEMA的PSNR值比傳統方法BM3D和WNNM提升約0.36～0.67 dB，比DnCN提升約0.20 dB，比FFDNet提升0.21 dB，比N3Net提升0.06 dB。在數據集BSD68與Urban100上，LEMA也均表現出較明顯的提升。