CAE-P:基于ADMM剪枝的圖像壓縮自編碼器

趙海萌

前言

圖像壓縮自編碼器CAE，亦即編碼器E及解碼器D，其訓練過程可以轉(zhuǎn)化為一個優(yōu)化問題，即對圖像失真率及圖像編碼比特數(shù)的最小化。有損圖像壓縮面臨著失真率及壓縮率的權(quán)衡問題，因此可以將上述優(yōu)化問題表述為：

minE，Dd+βR

其中d表示重構(gòu)圖像與原圖像之間的差距，R表示圖像編碼比特數(shù)，而β>0則控制上述兩個因素之間的平衡。解決這個優(yōu)化問題的過程中會遇到許多困難，其中最為重要的一個是如何表征圖像編碼比特數(shù)R。因此，本項目針對R的優(yōu)化，提出利用ADMM（AlternatingDirectionMethodofMultipliers）算法對CAE的表示層進行剪枝，即直接減小R，避開了額外訓練信息熵估計器的麻煩，并遵循訓練、剪枝、重訓練的順序，迭代地對CAE進行訓練（優(yōu)化d）和剪枝（優(yōu)化R），直至達到目標要求。

本文提出了CAE-P（CompressiveAutoEncoderwithPruning）模型，相較于現(xiàn)有的CAE模型，CAE-P模型顯得更為簡單直接，更易實現(xiàn)且參數(shù)量更小。實驗中，CAE-P模型在MS-SSIM（Multi-scaleStructuralSimilarityIndex）、SSIM等指標下均超越了現(xiàn)有的圖像壓縮算法。

CAE-P模型

一個基本的圖像壓縮自編碼器CAE由3部分組成：編碼器E、解碼器D和量化器Q。圖1為CAE-P的詳細模型架構(gòu)，“Convk/spP”表示卷積核大小為k×k、步長為s，并采用大小為P的鏡像Padding的卷積層，“ConvDown”表示將寬和高減半的卷積層。

在對編碼器E、解碼器D、量化器Q進行選擇及構(gòu)造，并解決優(yōu)化問題后，我們便得到了基于ADMM剪枝的圖像壓縮自編碼器CAE-P，它的基本結(jié)構(gòu)及運作方式如圖2所示。

原始圖像x經(jīng)過由卷積殘差塊構(gòu)成的編碼器E編碼，轉(zhuǎn)化為一組潛在表示形式z。在ADMM訓練階段，黑色和綠色通路激活，特征圖z經(jīng)量化器Q量化后得到，輸入同樣由卷積殘差塊構(gòu)成的解碼器D，解碼器D從中重構(gòu)出圖像，ADMM算法迭代地最小化重構(gòu)失真率d（x，），并對z剪枝，迫使其稀疏化。ADMM訓練階段結(jié)束后，綠色通路關(guān)閉，紅色通路開啟，進入重訓練階段。此時由于ADMM算法并不能保證z中僅剩余小于等于?個非零元素，只能使其他元素接近于零，因此需要對z中的近零元素強制剪除，再經(jīng)量化器Q量化后，得到最終編碼形式，并將此最終編碼形式輸入解碼器D中，得到重構(gòu)圖像。為了減小強制剪除對重構(gòu)圖像的影響，此時的重訓練便只需優(yōu)化minE，Dd（x，）。遵循訓練、剪枝、重訓練的步驟反復(fù)訓練，可以得到同時具有高圖像重構(gòu)質(zhì)量及高壓縮率的圖像編碼。這便是本文所提出的CAE-P模型的基本架構(gòu)及訓練方法。

圖2為CAE-P模型的訓練策略圖。其中，Encoder和Decoder均采用基于CNN的殘差塊構(gòu)建，Code為最終的壓縮后編碼。ADMM訓練過程中，黑色和綠色通路激活，ADMM算法在最小化重構(gòu)失真率的同時對z剪枝，迫使其稀疏化。ADMM訓練結(jié)束后，黑色和紅色通路激活，對z中的近零元素強制剪除，開始重訓練。

實驗過程

模型架構(gòu)

CAE-P模型由基于CNN和殘差塊的編碼器、解碼器組成。原始圖像首先經(jīng)過3次下采樣，每次包含下采樣卷積、批歸一化BN層和PReLU激活函數(shù);接下來經(jīng)過15層Bottleneck殘差塊，其中每次卷積后均連接BN層和PReLU;最后再經(jīng)過2層下采樣卷積，得到z。

接下來經(jīng)過量化器得到，量化器的梯度如下，設(shè)置為1。

量化后的編碼進入解碼器，解碼器的架構(gòu)與編碼器基本鏡像。編碼首先通過2層上卷積Sub-Pix層，每層包含卷積、BN層、PixelShuffle和PReLU激活函數(shù);然后經(jīng)過15層Bottleneck殘差塊，內(nèi)部結(jié)構(gòu)與編碼器中的相同;最后再經(jīng)過2層上卷積Sub-Pix層，最后通過Tanh激活函數(shù)，使激活值限制在（-1，1）之間，最后將其線性映射至（0，1）之間。

目標函數(shù)d的選擇

目前，常用的d的選擇包括MSE、MS-SSIM等。我們嘗試了多種設(shè)計后，最終選擇d（x，）=100·（1-MSSSIM（x，））+100·（1-SSIM（x，））+（45-PSNR（x，））+MSE（x，）。所有模型均以此為目標函數(shù)d訓練。

模型訓練

研究采用Adam作為優(yōu)化器，BatchSize設(shè)置為32，對ADMM優(yōu)化問題的第一步進行優(yōu)化。初始學習率設(shè)置為4·10-3，并隨著訓練過程動態(tài)衰減：每出現(xiàn)10個Epoch平均損失函數(shù)不下降，學習率減半。WeightDecay設(shè)置為一個較小值1·10-10。每經(jīng)過20個Epoch，執(zhí)行1次ADMM算法的第2步和第3步，第2步中的保留元素比例設(shè)置為10%。對于不同bpp的模型，修改編碼器的最后一層及解碼器的第1層的結(jié)構(gòu)，進行fine-tune。所有代碼在PyTorch框架下實現(xiàn)，每個模型在4塊NVIDIAGeForceGTX1080TiGPU上并行化訓練了300個Epoch，并在GitHub上開源。

數(shù)據(jù)集及其預(yù)處理

實驗選擇BSDS500作為訓練數(shù)據(jù)集，包含500張481×321的圖片。每次輸入網(wǎng)絡(luò)時圖像首先被隨機裁剪出一塊128×128大小的Crop，隨機進行水平和豎直翻轉(zhuǎn)，最后歸一化至（0，1）。對于測試集，實驗采用圖像壓縮領(lǐng)域內(nèi)通用的測試數(shù)據(jù)集KodakPhotoCD2，包含24張768×512的照片。

實驗結(jié)論

研究提出了CAE-P模型，避開了額外訓練熵編碼器的繁瑣與優(yōu)化的不直接，引入ADMM算法對編碼層剪枝，優(yōu)化更為直接，更易實現(xiàn)且參數(shù)量更小。在現(xiàn)有的使用熵編碼器的方法之外開辟并探索了一條新的道路：利用剪枝方法。實驗表明，實驗?zāi)Ｐ驮诙囗椫笜讼戮搅爽F(xiàn)有的圖像壓縮算法。

研究意義

本研究致力于改進圖像有損壓縮算法的壓縮率與圖像重構(gòu)質(zhì)量，首次將最先進的神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索技術(shù)遷移于其中，實現(xiàn)了超越現(xiàn)有圖像壓縮算法的性能，為圖像更高效地存儲與傳輸提供了基礎(chǔ)。

創(chuàng)新點

本研究提出了CAE-P模型，首次在圖像壓縮領(lǐng)域中引入剪枝方法，在使用熵估計器的方法之外開辟并探索了應(yīng)用架構(gòu)搜索方法（剪枝）的可能性，為未來引入更多更精巧的架構(gòu)搜索方法作了鋪墊。實驗表明，該模型在多項指標下均超越了現(xiàn)有的圖像壓縮算法，尤其是研究改進對象CAE。