圖鴨科技 CVPR 2019 CLIC 圖像壓縮挑戰(zhàn)賽冠軍方案全解讀

雷鋒網(wǎng) AI 科技評論按：在 CVPR 2019 CLIC 圖像壓縮挑戰(zhàn)賽中，圖鴨科技所提出的算法 TucodecSSIM 奪得了 MS-SSIM 和 MOS 兩項(xiàng)指標(biāo)的冠軍，算法 TucodecPSNR 奪得了 PSNR 指標(biāo)的冠軍，算法 TucodecPSNR40dB 則奪得高碼點(diǎn)圖像壓縮 Transparent Track 的冠軍。以下為圖鴨科技提供的技術(shù)解讀。

摘要：

近年來隨著人工智能技術(shù)的發(fā)展，基于深度學(xué)習(xí)的圖像壓縮技術(shù)已取得了飛速的發(fā)展。一個典型的基于深度學(xué)習(xí)的圖像壓縮框架包括：自編碼網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計(jì)、量化、碼率估計(jì)和率-失真優(yōu)化等幾個模塊。本文將主要介紹圖鴨科技在 CVPR 2019 CLIC 圖像壓縮挑戰(zhàn)賽上的相關(guān)技術(shù)方案，針對于比賽所設(shè)置的低碼點(diǎn)和高碼點(diǎn)壓縮兩個賽道，我們基于變分自編碼網(wǎng)絡(luò)設(shè)計(jì)了可進(jìn)行端到端優(yōu)化的圖像壓縮方案。該方案包括一個非線性編碼網(wǎng)絡(luò)、軟量化模塊、一個非線性解碼網(wǎng)絡(luò)和一個熵估計(jì)模塊。我們技術(shù)方案的特色之處總結(jié)如下：

1. 提出了基于全局特征分析的 non-lcao 注意力模塊，并融合進(jìn)編碼網(wǎng)絡(luò)和解碼網(wǎng)絡(luò)，以實(shí)現(xiàn)提升自適應(yīng)碼字分配性能的目的。

2. 設(shè)計(jì)了一種基于自適應(yīng)聚類的軟量化方法以降低量化損失。

3. 提出了能融合超先驗(yàn)子網(wǎng)絡(luò)和基于 pixel cnn++的上下文模型的碼率估計(jì)模塊。

得益于優(yōu)良的網(wǎng)絡(luò)結(jié)構(gòu)和算法設(shè)計(jì)，我們所提出的算法 TucodecSSIM 奪得了 MS-SSIM 和 MOS 兩項(xiàng)指標(biāo)的冠軍，算法 TucodecPSNR 奪得了 PSNR 指標(biāo)的冠軍，算法 TucodecPSNR40dB 則奪得高碼點(diǎn)圖像壓縮指標(biāo)的冠軍。接下來將具體介紹我們的算法方案：

方法介紹：

(1) 編碼網(wǎng)絡(luò)和解碼網(wǎng)絡(luò)

我們的主干壓縮自編碼網(wǎng)絡(luò)使用了如圖 1 所示的非對稱結(jié)構(gòu)，它包括卷積、非線性單元和殘差 non-local 注意力卷積等模塊。值得注意的是，通過使用殘差 non-local 注意力模塊來對特征的全局關(guān)聯(lián)性進(jìn)行捕捉和建模，圖像中的紋理、邊界等復(fù)雜部分能得到更好的重建。在 kodak 標(biāo)準(zhǔn)數(shù)據(jù)集上的實(shí)驗(yàn)表明，通過在編碼和解碼網(wǎng)絡(luò)結(jié)構(gòu)中融合 non-local 注意力模塊，能在 PSNR 指標(biāo)熵帶來 0.6db 的提升。

圖 1 編碼和解碼網(wǎng)絡(luò)結(jié)構(gòu)示意圖

(2) 量化

在現(xiàn)在的大部分方案中，取整量化是一種常用的方式。我們通過實(shí)驗(yàn)測評發(fā)現(xiàn)，這種直接將浮點(diǎn)數(shù)映射到整數(shù)的量化方式會極大的降低重建精度（PSNR 指標(biāo)至少降低 0.5db, MS-SSIM 指標(biāo)降低至少 1.5db）。為了降低量化帶來的精度損失，我們設(shè)計(jì)了一種基于自適應(yīng)聚類的軟量化方案，具體介紹如下：

給定可學(xué)習(xí)的中心點(diǎn)，可使用最近鄰分配的方式來計(jì)算量化值：

但式 (1) 的量化方式是不可導(dǎo)，因此將使用如下所定義的軟分配方式進(jìn)行替換，以保證在訓(xùn)練過程中能進(jìn)行端到端的優(yōu)化：

在用 tensorflow 進(jìn)行實(shí)現(xiàn)時，可以用如下的代碼對參數(shù)進(jìn)行量化：

(3) 先驗(yàn)概率和碼率估計(jì)

為了進(jìn)行碼率估計(jì)，使用拉普拉斯分布對壓縮特征的分布進(jìn)行表示，分布的參數(shù)包括均值和方差。為了對分布的均值和方差進(jìn)行端到端的計(jì)算，設(shè)計(jì)了如圖 2 所示的網(wǎng)絡(luò)結(jié)構(gòu)：

圖 2. 碼率估計(jì)模塊示意圖，該模塊包括超參自編碼網(wǎng)絡(luò)、上下文網(wǎng)絡(luò)和熵參數(shù)網(wǎng)絡(luò)三部分。

碼率估計(jì)模塊由三個子網(wǎng)絡(luò)構(gòu)成：超參網(wǎng)絡(luò)，參數(shù)為；上下文網(wǎng)絡(luò)，參數(shù)為；熵參數(shù)網(wǎng)絡(luò)，參數(shù)是。如圖 2 所示，超參網(wǎng)絡(luò)由超參編碼網(wǎng)絡(luò)、量化模塊和超參解碼網(wǎng)絡(luò)組成。超參網(wǎng)絡(luò)的量化特征也需要編碼，使用非參的概率密度估計(jì)方式進(jìn)行先驗(yàn)概率建模：

此外我們使用了 Pixelcnn++網(wǎng)絡(luò)結(jié)構(gòu)對特征的上下文關(guān)系進(jìn)行捕捉和建模；最后使用卷積模塊進(jìn)行熵參數(shù)網(wǎng)絡(luò)的構(gòu)建，并將超參網(wǎng)絡(luò)的輸出和上下文網(wǎng)絡(luò)的輸出進(jìn)行級聯(lián)來作為熵參數(shù)網(wǎng)絡(luò)的輸入來計(jì)算相應(yīng)的均值和方差。主干自編碼網(wǎng)絡(luò)壓縮特征的分布可以表示為：

最后碼率估計(jì)將由兩部分組成：一部分是對主干自編碼網(wǎng)絡(luò)中壓縮特征的估計(jì)碼率，一部分是超參自編碼網(wǎng)絡(luò)壓縮特征的估計(jì)碼率：

(4) 后處理

在方案 TucodecPSNR 中，我們使用了改進(jìn)的 266 算法作為基礎(chǔ)，但低碼率壓縮算法重建圖最顯著的缺點(diǎn)是存在偽影，并且很多紋理細(xì)節(jié)會丟失。為了改進(jìn)在低碼率條件下重建圖的質(zhì)量，我們設(shè)計(jì)一個有效地后處理模塊，后處理模塊的具體細(xì)節(jié)如圖 3 所示。

圖 3 后處理算法結(jié)構(gòu)示意圖

(5) 實(shí)驗(yàn)結(jié)果

我們從 CLIC 2019 訓(xùn)練集和 flickr.com 上收集了 5000 張高清圖片，并從中采集了百萬張的圖像塊作為訓(xùn)練集。為對網(wǎng)絡(luò)進(jìn)行高效訓(xùn)練，我們使用 tensorflow 平臺對相應(yīng)網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行實(shí)現(xiàn)。在方案 TucodecSSIM 中，進(jìn)行網(wǎng)絡(luò)訓(xùn)練使用的損失函數(shù)如下所示：

除上述損失函數(shù)外，也結(jié)合了對抗生成網(wǎng)絡(luò)對壓縮網(wǎng)絡(luò)進(jìn)行端到端的訓(xùn)練。為了滿足比賽的 0.15bpp 的約束要求，我們訓(xùn)練了 5 個模型分別對應(yīng)λ=0.2/0.3/0.4/0.5/0.6，最后使用動態(tài)規(guī)劃算法進(jìn)行碼字分配。

在方案 TucodecPSNR40db 中，進(jìn)行網(wǎng)絡(luò)訓(xùn)練所使用的損失函數(shù)如下所示：

在模型訓(xùn)練完畢后，為了滿足比賽的 PSNR 值不小于 40db 和 MS-SSIM 值不小于 0.993 的約束，共訓(xùn)練了 5 個模型，對應(yīng)的 λ=4096/4800/5500/6500/8000。最后使用動態(tài)規(guī)劃算法進(jìn)行碼字分配。在方案 TucodecPSNR 中，使用了我們改進(jìn)的 H266 算法作為基礎(chǔ)，并結(jié)合后處理網(wǎng)絡(luò)進(jìn)行性能的提升，我們給出了三個模型，分別對應(yīng) QP 36/37/38，并最后進(jìn)行碼字分配以滿足 0.15bpp 約束的要求。所提出算法的測評指標(biāo)如下表所示：