深度學習輕量化偵察圖像壓縮網絡

諶 宇,諶德榮,崇魁奇,王澤鵬,張 凱

(1.北京理工大學,北京 100081;2.北京航宇天穹科技有限公司,北京 100043)

0 引言

小型偵察平臺具有靈活性高、隱蔽性強等特點,是快速獲取偵察圖像的重要技術手段。為了保證獲取圖像的時效性,小型偵察平臺需要將獲取的偵察圖像壓縮后通過無線信道傳回指揮中心,圖像壓縮性能是決定獲取圖像質量的關鍵。

根據處理方法的不同,圖像壓縮算法可以分為傳統壓縮算法和基于深度學習的壓縮算法兩類。

傳統壓縮算法主要通過預測、變換去除圖像的空間冗余,通過熵編碼去除統計冗余,進而達到壓縮數據的目的。例如,JPEG2000[1]壓縮算法通過小波變換來去除圖像中的空間冗余,利用算數編碼來去除統計冗余,是傳統圖像壓縮算法中性能最好的。

基于深度學習的圖像壓縮算法主要通過神經網絡來去除圖像中的冗余信息。根據使用的網絡模型結構不同,基于深度學習的圖像壓縮方法可以分為基于遞歸神經網絡的、基于變分自編碼器的和基于生成對抗網絡的圖像壓縮算法三類。

基于遞歸神經網絡(RNN)的圖像壓縮算法通過迭代調用基于RNN的編碼器來壓縮圖像或圖像殘差的數據量,并利用熵編碼進一步降低數據量。雖然能精確地控制壓縮比,但是存在網絡模型復雜度較高、計算量大、圖像壓縮質量不佳等問題。

基于變分自編碼器(VAE)的圖像壓縮算法利用卷積神經網絡搭建由編碼器和解碼器組成的端到端學習網絡。此算法可以超過JPEG2000的圖像壓縮質量,但是單個網絡模型只能實現一種壓縮比的圖像壓縮,而且使用了熵編碼處理,不利于硬件的并行加速處理。

基于生成對抗網絡(GAN)的圖像壓縮算法在基于VAE的端到端圖像壓縮模型的基礎上引入判別器以構成GAN,并且以金字塔式結構設計了編碼器和判別器,在去除圖像冗余的同時保留更多細節信息,提高了圖像質量,但同樣存在模型復雜、計算量大等問題。

為滿足小型偵察平臺對圖像壓縮算法計算復雜度低、計算量小、模型參數少等要求,提出基于深度學習輕量化偵察圖像壓縮網絡,通過輕量化編碼網絡設計、高質量解碼網絡設計和損失函數設計以降低編碼網絡的復雜度和提升圖像壓縮質量。該網絡在編碼端采用三個卷積模塊將圖像直接映射為服從均勻分布的二進制碼流,得到壓縮后數據,進而去除了不利于硬件并行處理的熵編碼操作;為了降低卷積模塊的計算量和參數量,采用深度可分離卷積代替傳統卷積,并利用分組卷積+通道重排優化卷積模塊;在解碼端通過轉置卷積和殘差模塊提升特征的提取和恢復能力,進而提高解碼圖像質量。為了使壓縮后的數據服從均勻分布使其包含的信息量最大化,并使圖像重建誤差最小化,設計了能夠同時反映壓縮后數據信息熵和重建誤差的目標函數。

1 卷積神經網絡

卷積神經網絡(CNN)是一種包含卷積計算且具有深層次結構的神經網絡,通常由數據輸入層、卷積層、激活函數層、池化層與全連接層組成。

數據輸入層一般是對輸入卷積網絡的數據進行預處理,例如去均值、歸一化等。該層的主要目的是加快訓練過程中網絡的收斂速度,調整輸入數據的取值范圍,防止超過激活函數的值域限制。

卷積層主要通過卷積核完成對數據的卷積操作。該層的目的是獲取數據更深層次的特征信息。在卷積神經網絡中,決定某一層輸出結果中一個元素所對應的輸入層的區域大小,被稱作感受野。卷積層的層數越多,感受野就越大;感受野越大,網絡能提取的信息就越充分。

激活函數的主要作用是向網絡中引入非線性,加入激活函數之后,深度神經網絡具備了非線性映射學習能力。

全連接層一般處在網絡的末端,在整個卷積神經網絡中起到“分類器”的作用,將網絡中生成的各種特征整合到一起。 該層的參數較多,容易出現參數冗余,降低訓練速度等問題。

2 輕量化圖像壓縮網絡

通過輕量化編碼網絡設計、高質量解碼網絡設計和損失函數設計以降低編碼網絡的參數量與計算量和提升圖像壓縮的質量。

2.1 編碼網絡設計

編碼網絡的設計以編碼結構簡單、編碼參數少、編碼計算量少為指導思想。編碼網絡使用卷積神經網絡將圖像直接映射為二進制碼流,同時為了減少編碼的計算量,引入深度可分離卷積與分組卷積,大大降低了網絡的參數量與計算量。編碼網絡結構如圖1所示。

圖1 編碼網絡結構示意圖Fig.1 Schematic diagram of coding network structure

編碼網絡主要包含三個模塊:

第一個模塊包含一個下采樣卷積層和一個ReLU激活函數。下采樣卷積層對輸入數據進行步長為2的下采樣卷積處理。該層卷積核尺寸為3×3,卷積核個數為32個。

第二個模塊采用了深度可分離卷積以及分組卷積+通道重排。該模塊包含一個ReLU激活函數、一個下采樣深度卷積層和一個點卷積層。下采樣深度卷積層,進行步長為2的下采樣,其卷積核為尺寸3×3,個數為32個。點卷積層根據分組卷積原理將該層分為兩個部分,每部分卷積核尺寸為1×1×16,卷積核個數為16個。

第三個模塊也采用了深度可分離卷積。該模塊也包含一個下采樣深度卷積層,以及一個點卷積層。下采樣深度卷積層采樣步長為2,卷積核尺寸為3×3,卷積核個數為32個。點卷積層卷積核尺寸為1×1×32,該層卷積核數量與壓縮比有關。

經過三個模塊處理后,輸出數據經二進制量化得到壓縮后的數據。二進制量化處理將網絡輸出的數據量化為位寬為1 b的0或1,進一步提高壓縮比。在測試過程中二值量化采用以0.5為分界的階躍函數;在訓練過程中為了保證梯度能夠反傳,采用Sigmoid函數代替階躍函數完成量化。

壓縮比與第三模塊點卷積層卷積核的數量相關。設編碼網絡輸入的灰度圖像長和寬分別為H,W,第三模塊點卷積層的卷積核個數為M,則編碼網絡輸入的數據量為H×W×1,編碼網絡輸出數據量為H/8×W/8×M,此時數據量壓縮比為64/M。原始圖像像素位寬為8 b,二值量化后數據位寬為1 b,因此最終壓縮比為64/M×8。當輸入圖像尺寸為128×128,壓縮比設定為64時,第三模塊點卷積層個數為8,編碼端參數量為1 632,編碼器浮點運算量(FLOPs)為4 069 376(一次乘加運算FLOPs為2)。

2.1.1深度可分離卷積

在編碼網絡中,第一模塊輸入通道數為1,第二、第三模塊輸入通道數為32。計算量主要集中在第二模塊和第三模塊。在這兩個模塊中使用深度可分離卷積,可以大大減少編碼網絡的參數量與計算量[2-3]。

深度可分離卷積包括深度卷積(DW)和點卷積(PW)兩個部分。相比常規的卷積操作,其參數量和計算量比較低。

在深度卷積中,每個輸入通道的特征圖分別與一個卷積核進行卷積,而非每個輸入通道的特征圖與所有卷積核進行卷積。深度卷積的計算結果缺少不同通道在相同空間上的信息,通常需要在其后加一個點卷積,以融合不同通道之間的信息。深度卷積的卷積核數與輸入通道數一致,不會改變數據的維度。

點卷積的卷積核尺寸為1×1×M,M為輸入通道數。點卷積運算主要是對每個輸入特征圖的像素點在通道的維度上進行線性組合。

若編碼網絡第二、第三模塊不使用深度可分離卷積,則編碼網絡結構如圖2所示。

圖2 剝離深度可分離卷積的編碼網絡結構Fig.2 Schematic diagram of code network structure without using depth-separable convolution

圖2中,第二模塊只保留了一層下采樣卷積層和一個ReLU激活函數。卷積層分為兩個部分,每個部分卷積核尺寸為3×3×16,卷積核個數為16個;第三個模塊也只保留了一層下采樣卷積層,該層卷積核尺寸為3×3×32,卷積核個數為8個。此時編碼網絡參數量為7 200,是使用深度可分離卷積的編碼網絡參數量的4.4倍;在輸入圖像尺寸為128×128的情況下,浮點運算量為11 750 400,是使用深度可分離卷積的編碼網絡運算量的2.9倍。

2.1.2分組卷積+通道重排

在編碼網絡三個模塊中第二模塊的輸入通道與卷積核的數量最多,可通過對該模塊采用分組卷積[4-5]+通道重排[6-8]的方法以減少參數量及計算量。

分組卷積是將輸入通道以及卷積核分為T組,組內進行常規卷積,組與組之間沒有聯系。其中,每一組輸入通道和卷積核個數為原來的1/T。

分組卷積在降低計算量的同時削弱了通道之間的相關性,引入ShuffleNet算法的通道重排模塊可提高對通道相關性利用。分組卷積與通道重排結構如圖3所示。

圖3 分組卷積與通道重排示意圖Fig.3 Schematic diagram of group convolution and channel shuffle

在編碼網絡的第二模塊中,點卷積層使用了分組卷積,將該層卷積分兩組,每組由16個1×1×16卷積構成,輸入通道數為第二模塊深度卷積輸出通道數的一半,即每一組的輸入通道數為16。

為了提高對通道相關性的利用,參考ShuffleNet算法,首先將第一模塊下采樣卷積層輸出的32個通道進行分組,通道1～16為一組,通道17～32為一組;然后將兩組通道進行交錯排列,并將重排后的特征圖傳入下一個處理網絡。

若編碼網絡第二模塊不使用分組卷積+通道重排,則編碼網絡結構如圖4所示。

圖4 剝離分組卷積+通道重排的編碼網絡結構圖Fig.4 Schematic diagram of code network structure without using Group convolution and channel shuffle

與圖1相比,圖4中第一、第三模塊未發生變化,第二模塊的點卷積層卷積尺寸變為了1×1×32,卷積核個數變為32個。此時編碼網絡參數量為2 144,是使用分組卷積+通道重排編碼網絡參數量的1.3倍;在輸入圖像尺寸為128×128的情況下,編碼網絡浮點運算量為5 117 952,是使用分組卷積+通道重排編碼網絡參數量的1.25倍。

表1、表2比較了在128×128分辨率下本文設計的編碼網絡、不使用深度可分離卷積的編碼網絡、不使用卷積分組+通道重排的編碼網絡,三種編碼網絡的參數量與計算量。參數量見表1,計算量見表2。

表1 編碼網絡參數量對比表Tab.1 Code network parameter number comparison table

表2 編碼網絡計算量對比表Tab.2 Comparison table of computations at coding network

由表1、表2可知,深度可分離卷積和卷積分組+通道重排可以使編碼網絡參數量、計算量大大降低。

2.2 解碼網絡設計

解碼網絡以提高重建圖像質量為設計目的。提出的解碼網絡結構如圖5所示。

圖5 解碼網絡結構示意圖Fig.5 Decoding network structure diagram

在圖5中,解碼網絡主要由3種模塊組成:卷積層、轉置卷積、殘差模塊。

在解碼網絡中,第一個卷積層尺寸為3×3,卷積核個數為64個;最后一個卷積層尺寸為3×3,卷積核個數為1個;其余的卷積層卷積尺寸為3×3,卷積核個數為32個。

轉置卷積用來恢復圖像分辨率。轉置卷積的核長為4,步長為2,卷積核個數為32個。這樣可以避免由棋盤效應帶來的重建圖像質量下降的問題。

殘差模塊可以緩解網絡過深造成的網絡退化和梯度消失問題,有利于提升梯度傳遞的效率。殘差模塊結構中包含兩層卷積層和兩個ReLU函數。

2.3 損失函數設計

(1)

式中:壓縮量y的熵值如式(2)所示。

(2)

式中:n為壓縮量y的取值范圍,pi為對應y值出現的概率。

(3)

3 網絡訓練與測試

網絡訓練所用平臺配置如下:CPU為Intel Xeon(R)Silver 4210R @2.40 GHz×40,GPU為 NVIDIA RTX 3090 TI,內存為64 GB。

3.1 數據集準備

實驗數據來自2018和2019屆CLIC圖像壓縮挑戰賽,共616張圖像。其中隨機選擇587張原始圖像隨機裁剪為尺寸為256×256的圖像,共獲得24 078張圖像用于構成訓練集。剩余42張原始圖像,裁剪為6 692張大小為128×128的圖像作為驗證集。選用Set5,Set14[9],SunHays80,BSD100,Urban100[10]作為本次實驗的測試集,將測試集圖像裁剪為128×128的圖像進行測試。

3.2 模型訓練

使用Adam優化器進行訓練,初始學習率設置為0.000 1,在達到第80輪的時候學習率降低10倍,共訓練120個輪。模型效果的評估準則為PSNR和SSIM[11],PSNR計算公式為

(4)

式中:(I,K)MSE表示圖像I和K殘差值的平方。

(5)

SSIM計算公式為

(I,K)SSIM=L(I,K)·C(I,K)·S(I,K),

(6)

(7)

(8)

(9)

式中,uI和uK分別表示圖像I和K的均值;σI和σK分別表示圖像I和K的標準差;σIK表示圖像I和K的協方差;C1,C2,C3為常數,避免分母為0。

3.3 測試結果

對于測試集,輕量化偵察圖像壓縮網絡與JPGE2000算法編碼所需時間如表3所示。

表3 編碼時間對比表Tab.3 Code time comparison table

由表3可知,本文方法編碼速度約是JPGE2000算法編碼速度的11倍。

JPGE2000算法和本文算法重建圖像質量如表4所示。選用5個測試集PSNR與SSIM的加權平均作為重建質量的度量。

表4 重建質量對比Tab.4 Reconstruction quality comparison

由表4可知,輕量化偵察圖像壓縮網絡與JPGE2000算法相比,PSNR提高了3.85 dB,SSIM提高了0.27,重建質量遠超于JPGE2000算法。解碼端使用轉置卷積能獲得更好的重建圖像質量。

圖6給出了JPGE2000和本文算法重建偵察圖像(分辨率為128×128)的結果。進一步說明了輕量化偵察圖像壓縮網絡的重建質量高于JPGE2000。

圖6 128×128分辨率下圖像對比Fig.6 Image comparison at 128×128 resolution

為了驗證深度可分離卷積、分組卷積+通道重排以及轉置卷積的作用,本文進行了剝離實驗以驗證其效果。對比了用普通卷積替換深度可分離卷積、普通卷積替換分組卷積、最近鄰插值替換轉置卷積三種情況與本文方法的重建圖像精度,結果如表5所示。

表5 重建質量對比Tab.5 Reconstruction quality comparison

對比表1、表2以及表5中無深度可分離卷積、無分組重排以及輕量化偵察圖像壓縮網絡可知,引入深度可分離卷積與分組卷積+通道重排不僅可以減少編碼網絡的參數量與計算量,還能提高重建圖像的質量。用轉置卷積進行上采樣比用最近鄰插值進行上采樣能獲得更好的重建圖像質量。

4 結論

本文設計了一種基于深度學習輕量化偵察圖像壓縮網絡,編碼網絡由三個卷積模塊組成,并通過深度可分離卷積、分組卷積+通道重排的方式降低編碼網絡的參數量與計算量。該編碼網絡能夠將輸入數據直接映射為服從均勻分布的二進制碼流,從而完成圖像壓縮;解碼端使用轉置卷積、殘差模塊增加特征的提取與重建能力。實驗結果表明,與JPGE2000算法相比,基于深度學習輕量化偵察圖像壓縮網絡的PSNR提高了3.85 dB,SSIM提高了0.27,編碼時間降低91%。