「JPEGAI」End-to-end optimized image compression

語言: CN / TW / HK

本文已參與「新人創作禮」活動,一起開啟掘金創作之路。

論文地址:https://arxiv.org/abs/1611.01704

代碼地址:https://github.com/tensorflow/compression

Balle  google大神,入門必讀。該文章僅供本人筆記用,如果問題歡迎討論。

簡介

整體算法分為三個部分:非線性分析變換(編碼器),均勻量化器和非線性合成邊變換(解碼器),在整個訓練框架上聯合優化整個模型的率失真性能。在某些條件下,鬆弛損失函數可以解釋為由變分自動編碼器實現的生成模型的對數似然,即和變分自編碼器的損失函數某種情況下有點像。另外,每一個折衷參數對應一個壓縮模型,所以整個R-D曲線是由模型構成的R-D點形成的。

內容

整體算法結構以及流程如下:

!](https://img-blog.csdnimg.cn/20201205140711473.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQyMjgxNDI1,size_16,color_FFFFFF,t_70)

在這裏插入圖片描述 $x$ 與 $\hat{x}$ 分別代表輸入的原圖和經過編解碼器後的重建圖片。 $g_a$表示編碼器提供的非線性分析變換,$y$即由輸入圖片經過編碼器網絡後得到的潛在特徵,通過量化器$q$後,得到 量化後結果:$\hat{y}$,再通過 $g_s$解碼器重建圖片結果.

其中,通過對$y$的碼率估計得到 $R$,計算原圖 $x$ 和 $\hat{x}$ 的失真得到$D$,$g_p$是指一種失真變換,可以將原圖和重建圖進行通過 $g_p$ 轉化到感知空間上進行計算,,直觀的作用可以理解為失真計算:$D=d(x,\hat{x})$,例如PSNR,MS-SSIM,或者其他感知域如 VMAF 等。通過得到的 $R$ 和 $D$ 進行率失真聯合優化,定義損失函數為: $$L=\lambda\cdot D+R$$通過 $\lambda$ 參數進行碼率的選擇控制, $\lambda$ 大則訓練得到的模型的重建圖失真,而壓縮後的碼流大,反之亦然。

激活層改進:受到生物神經元的啟發,在編碼器與解碼器中,論文使用了GDN(廣義除數歸一化),該激活層在論文中《Density modeling of image using a generalized normalization transformation》有着詳細的描述,該激活層實踐中比BN層更適用於圖像重建工作。其中與BN層的差別:BN層重新調整網絡中線性濾波器的響應,以便將其保持在合理的操作範圍內。這種類型的歸一化不同於局部增益控制,因為在所有空間位置上重新縮放因子是相同的。此外,一旦訓練完成,縮放參數通常是固定的,這將歸一化轉換為數據的仿射變換,不同於GDN那樣,GDN是空間自適應的,並且可能是高度非線性的。

量化問題:傳統的圖像編碼中,量化有着不可微分的問題,為分段函數,且在非邊界點的導數為0,截斷反向傳播,在邊界點則不存在導數。作者採用基於概率模型連續鬆弛的代理損失函數(沒理解什麼意思),用加性均勻噪聲代替量化步驟(就是加了-0.5到0.5區間的均勻噪聲)。

前向傳播

在非線性分析變換(編碼器)中 $g_a$ 包含有三個階段:卷積,降採樣(就是用的卷積層進行的降採樣工作),GDN歸一化。以下公式預警:

卷積

$$v_{i}^{k}(m,n)=\sum\limits_{j}(h_{k,i j}*u_j^{k})(m,n)+c_{k,i}$$

其中 $v_{i}^{k}(m,n)$ 是$i$個輸出的通道在第$k$層的網絡的輸入,總共有3層,所以k最大為3,其中($m$,$n$)代表長寬的二維的位置。$*$是卷積操作,上述公式描述輸入在 經過卷積層後的輸出。

採樣

$$ w_i^{k}(m,n)=v_{i}^{k}(s_km,s_kn)$$

上述公式,$s_k$是樣因子,即幾倍採樣。一般採樣與卷積運算在代碼中統一實現。

GDN $$u_i^{k+1}(m,n)= \dfrac{w_i^{k}(m,n)}{\sqrt{(\beta_{k,i}+\sum\limits_jr_{k,ij}(w_j^{k}(m,n))^2)}}$$

所有的待優化參數$h$(卷積核權重),$c$(卷積核偏置),$r$(GDN的歸一化係數)都在整個端到端中被優化。

後續在解碼端中的公式不例舉遼,這部分內容説白了就是把神經網絡作用給抽象公式化。

變換優化問題

     在傳統的圖像編碼中,對於DCT生成的係數需要進行量化,通常採用了矢量量化的形式並且結合熵編碼進行率失真控制。在該論文並沒有直接在編碼空間進行最優量化(矢量),而是通過一個固定的標量量化器進行量化(四捨五入),試圖讓分線性變換(編碼器)進行自動學習,以學習到的方式扭曲量化空間,從而有效實現矢量量化。整體框架通過以下公式進行優化:$$L=-E[log_2P_q]+\lambda\cdot E[d(z,\hat{z})]$$通過計算熵與失真進行率失真優化,其中兩個期望值通過訓練集的平均值進行擬合。並且通過非線性分析學習量化特徵即矢量量化結構,則本文采用標準量化:$$\hat{y_i}=q_i=round(y_i)$$,$i$ 覆蓋了所有的,待編碼的數值,並且 $\hat{y_i}$ 的邊緣密度一系列離散概率密度質量給出:$$P_{q_i}(n)=\int_{n-1/2}^{n+1/2}p_{y_i}(t)dt $$即$p_{y_i}(t)$是$y$的概率密度函數,由於四捨五入,在某一整數的(-0.5,0.5)區間內都會量化為該整數,則通過積分的形式計算這一區間內的數值的出現概率,得到量化後的整數的出現概率。

     上述公式均涉及到量化問題,但是量化會導致不可微分,阻斷反向傳播優化的問題。論文中採用了添加(-0.5,0.5)範圍的均勻噪聲,在訓練過程中,即採用這種形式近似可微以用於反向傳播的優化,在推理過程中,則依舊使用 $round$ 函數進行四捨五入(因為不用進行優化了),使用均勻噪聲有兩個優點:

首先,$\tilde{y} = y+noise$ 的密度函數是$q$的概率質量函數的連續鬆弛

<div align=center>

$y_i$是編碼空間的元素,就是需要被編碼的值,$\hat{y_i}$ 是四捨五入後的值,$\tilde{y_i}$ 是通過添加噪聲後的值,上述是三者的概率密度函數(PDF),離散的$\hat{y_i}$是概率質量函數(PMF),其中黑點的數值由實線在各自區間內取積分得到的,但是其實不一定等於虛線在區間內的積分,只能是近似關係(論文中説相等,我不信=-=)。即可以通過這種近似關係,合理等於兩者的微分熵,換個説法就是加均勻噪聲後,不影響碼率的估計。

其次,獨立的均勻噪聲就其邊際矩(我也不知道什麼是邊際矩)而言近似於量化誤差,並且經常被用作量化誤差的模型(Gray and Neuhoff,1998)。

針對y的寬鬆概率模型和熵代碼在代碼空間中假設獨立的邊際,則通過進行參數建模以減少碼率估計模型與實際模型之間的差距,具體使用精細採樣的分段線性函數,這些函數與一維直方圖的更新類似(請參見附錄)即通過神經網絡生成編碼點的PMF,通過PMF的積分求得每個特徵點對應的可導形式的概率值,通過對概率值求自信息得到估計的碼率點,在前向傳播中,通過訓練階段生成的神經網絡生成分位點與CDF,將CDF與待編碼點輸出range coder進行熵編碼。 以用於熵編碼。 由於編碼數值的概率密度函數被均勻噪聲平滑,這部分會導致一定的誤差,但是通過減小採樣間隔可以使得模型誤差任意小。

整個過程可以公式得到:$$L(\theta,\phi)=E_{x,\triangle y}[-\sum_ilog_2p_{\bar{y_i}}(g_a(x;\phi)+\triangle y);\psi_{(i)}+\lambda d(g_p(g_s(g_a(x); \phi)),g_p(x))]$$量化以及$\psi$的分段線性逼近很適合隨機優化(我看不出=-=!)

變分推導問題

再論,等下一篇超先驗網絡結構模型進行推導吧。

實驗結果

在這裏插入圖片描述 左圖為量化所引起的誤差和均勻噪聲引起的誤差的關係,右圖為離散情況下的熵率以及可微分情況下的熵率的關係,可以得出,本文提出的量化模型以及熵率估計模型與實際中量化引起的誤差和待編碼值得熵率近似一致。

在這裏插入圖片描述 以及在實際中,得到得R-D性能,PSNR,MS-SSIM性能指標熵超過了JPEG2000壓縮標準。其中在PSNR度量標準上是通過MSE為失真指標進行訓練的,而MS-SSIM是在MS-SSIM指標上進行訓練到。並不是一種失真模型對應兩個失真指標。即MSE訓練了六個lambda的模型,MS-SSIM訓練了六個lambda的模型。 在這裏插入圖片描述 在這裏插入圖片描述 在這裏插入圖片描述 通過主觀感受可以看出,本文提出的算法優於JPEG以及JPEG2000。

結論

該論文還是着重與量化和熵率估計的問題,論述了添加均勻噪聲的量化的合理性以及對於熵率估計的影響,包括提出瞭如何使得熵率估計模型可微分以及CDF的學習以用於熵編碼的情況。並且確定了在圖像壓縮框架中,標量量化可以通過編碼器的非線性分析達到矢量量化的效果(期望這種非線性分析可以學習得到)。