端到端图像压缩《Asymmetric Gained Deep Image Compression With Continuous Rate Adaptation》

论文地址：https://openaccess.thecvf.com/content/CVPR2021/papers/Cui_Asymmetric_Gained_Deep_Image_Compression_With_Continuous_Rate_Adaptation_CVPR_2021_paper.pdf

代码地址：https://github.com/mmSir/GainedVAE

华为可变速率文章，代码非官方部署，一个小哥实现的，该文章仅供本人笔记用，如果问题欢迎讨论。

一 简介

       提出了一种连续速率可调的学习图像压缩框架，即非对称增益变分自动编码器（AGVAE）。 AG-VAE 利用一对增益单元在一个模型中实现离散速率自适应，而额外的计算可以忽略不计。并且通过使用指数插值，在不影响性能的情况下实现连续速率自适应。除此之外，部署了非对称的熵模型提高模型的性能增益。

二 内容

2.1 目前方法的缺陷

       在之前的图像压缩工作中，都是通过改变损失函数中的 $\lambda$ 超参数调整模型的码率，这会导致无法将图像压缩至固定码率点的情况，并且要为此训练多个模型，需要消耗大量的训练时间和存储模型所需要的空间。为此，单个模型能够覆盖多个码率的情形有很大的应用需求。
       在以往的可变速率技术方案中，基于RNN方案进行渐进式的图像编码，但是RD性能比较差，基于条件的conditional卷积网络复杂度高并且占用内存大，可变的量化bin size方式会导致性能的下降，此外对于BottleNeck层的尺度缩放方案在低码率的情况下会掉性能。

       在编解码方案中，不同通道对最后的重建质量的影响是不同的。作者探索基线方案中，被量化中的前32个通道信息对最后重建图像质量的影响，得出不同通道有不同重要性的结论，并且对通道进行scale的缩放，被量化后的潜在表示值乘以尺度缩放因子，得到潜在表示缩放后的重建质量。

2.2 整体方案

       整体方案还是和Google[1]的网络一致，对比框架优化了 Gain Unit 单元，扩展了自回归模型中的Mask Convolution，从1个5x5和扩展成 3x3，5x5，7x7的网络，并且文章中有优化熵模型，从单高斯模型扩展到高斯分布的两侧采用不同方差的半边高斯分布。

图1.整体网络框架图

2.2 Gain Unit

       标记编码器的输出$y\in R^{c,h,w}$即有c个通道（一般为192）w,h的宽高。$y_i\in R^{h,w}$则表示单个通道的潜在表示，其中$i\in C$，对应的Gain Unit 单元是有一个矩阵$M\in R^{c,n}$，表示这个矩阵实际上是为每一个通道的潜在表示分配一个长度为 n 的向量，$m_s\in m_{(s,0)},m_{(s,1)},m_{(s,2)},...m_{(s,c-1)}$。即每个$m_s$ 是一个长度为n的向量，对于每个通道的操作表示如下：$\overline{y}=y_i\times m_{s,i}$ 这里就是关于 Gain Unit的说明，每一个通道上的潜在表示都会乘以对应向量中的某个值，下面介绍基于Gain Unit的离散可变速率框架：

       编码器的输出 $y$ 经过Unit Gain单元进行处理缩放之后，得到 $\overline{y}$，并且需要经过量化，得到量化后的潜在表示$\hat{y}=round(\overline{y})$，解码端会同样部署Inverce-Unit Gain，从熵解码器中得到$\hat{y}$ ，然后进行对应拟变换得到$y^{{}^{\prime }}=InverseGain(\widehat{y_i}\times m_{s,i}^{{}^{\prime }})$。

       整体框架的损失函数优化主流的损失函数基本保持一致：

其中，$R_{\phi }$ 项表示码率， $D$ 表示失真， ${\beta }_s$ 表示训练模型中，失真和码率的权衡，${\beta }_s$ 越大，则表示模型越注重重建图像的质量，${\beta }_s$ 是从一组预定好的参数集中选取的。有${\beta }_s\in B$ ，其中B的长度为n，定义好了一系列的权重值。此外，可以发现， ${\beta }_s$ 的长度为n，而对于每一个Unit Gain 矩阵，每一个通道的scale向量的长度也是n。不同与之前的模型，每一个模型的训练 ${\beta }_s$ 为预定的单个数值，而此处的 ${\beta }_s$ 在训练中从 B 的预定义池中随机选取的，并且在选取 ${\beta }_s$后，得到对应的 $m_s,m_s^{{}^{\prime }}$。

       在推理阶段，可以训练的Unit Gain 矩阵中获取到有映射关系的$m_s,m_s^{{}^{\prime }}$ 对潜在表示$y$和$\hat{y}$进行缩放，得到对应几个离散情况下的离散点，如下图所示，训练了基于mse loss和(1-msssim) loss的两个模型。并且通过修改Gain Unit矩阵中的对应$m_s,m_s^{{}^{\prime }}$ 向量，得到的离散RD曲线，离散的模型记为DVR模型。

2.3 连续可变速率模型

本文通过了采用不同的 $m_s,m_s^{{}^{\prime }}$ 矢量对完成单模型多码率的模型设置，与此同时，可以对 $m_s,m_s^{{}^{\prime }}$ 和 $m_{s-1},m_{s-1}^{{}^{\prime }}$ 进行差值完成连续可变速率的实现。为了确保不同的$m_s,m_s^{{}^{\prime }}$ 之间的对于潜在表示$y$和$\hat{y}$ 的缩放结果是一致的，对不同的 $m_s,m_s^{{}^{\prime }}$ 有以下约束：
$$m_s\ast m_s^{{}^{\prime }}=m_t\ast m_t^{{}^{\prime }}=C$$
$m_s,m_s^{{}^{\prime }}$ 和 $m_t,m_t^{{}^{\prime }}$ （$r,t\in [0,1,...n-1]$）表示不同的增益矢量单元对应在不同的${\beta }_s$ 和 ${\beta }_t$。有以下公式：

此处 $m_v,m_v^{{}^{\prime }}$ 表示 $m_r,m_r^{{}^{\prime }}$ 和 $m_t{m}_t^{{}^{\prime }}$ 之间的差值系数，通过控制参数$l$来表示 $m_v,m_v^{{}^{\prime }}$ 的取值情况，当$l$从0取到1时，模型能够取到两个离散点$m_r,m_r^{{}^{\prime }}$ 和 $m_t{m}_t^{{}^{\prime }}$之间所有的连续的码率点，从而实现连续可变速率的目的。结果如下图所示：

可以从上图看出，插值后的模型RD性能基本无损甚至比拟合的曲线RD性能更高。

2.4 高斯熵模型

       在之前的工作中，文章大致采用了三种概率密度函数对参数进行建模：单高斯概率密度函数，混合高斯概率密度函数，单拉普拉斯概率密度函数。其中单高斯概率密度函数表达如下：

$\mu$和 $\sigma$表示概率密度函数的均值和方差，上述对称的方式对于建模的精准度不够高，不够freedom ，因此使用一种非对称的熵模型，具体如下：

公式中，$\mu$还是表示原来的含义，但是对于均值左右两侧的方差采用非对称的形式，即左边和右边采用不同的方差，${\sigma }_l和{\sigma }_r$则分别表示两侧的方差参数。

2.5 实验补充说明

在实验中，对应训练msssim的模型使用的 B m s s s i m = { 0.07 , 0.03 , 0.007 , 0.003 , 0.001 , 0.0006 } B_{msssim}={\{0.07,0.03,0.007,0.003,0.001,0.0006\}} Bmsssim​={0.07,0.03,0.007,0.003,0.001,0.0006} 六个数值，和训练psnr指标的模型 B m s e = { 0.05 , 0.03 , 0.007 , 0.003 , 0.001 , 0.0003 } B_{mse}={\{0.05,0.03,0.007,0.003,0.001,0.0003\}} Bmse​={0.05,0.03,0.007,0.003,0.001,0.0003}，以训练psnr模型为例，在训练的每个batch，会随机生成一个索引值 s ∈ { 0 , 1 , 2 , 3 , 4 , 5 } s\in{\{0,1,2,3,4,5\}} s∈{0,1,2,3,4,5},对应能够取到 B m s e = { 0.05 , 0.03 , 0.007 , 0.003 , 0.001 , 0.0003 } B_{mse}={\{0.05,0.03,0.007,0.003,0.001,0.0003\}} Bmse​={0.05,0.03,0.007,0.003,0.001,0.0003}中的对应值，与此同时为整个Unit Gain单元初始化一个长度矩阵 $M[6][192]$的矩阵，同样Inverce Gain也会有一个同样的矩阵，上面的$m_s$就是表示， 对应可以通过s索引得到$m_s=M[s][:]$，在训练开始前，可以把$M$矩阵初始化为常量，在每个batch的迭代中，从M矩阵中提取的$m_s$都会被更新，并且保存。

三 性能结论

整体的性能是基于图一框架以及对应的使用非对称熵模型得到的，除此之外，还使用了$论{文}^1$中的attention机制和$论{文}^2$中Universal量化技术和$论{文}^3$中采用的并行上下文技术，整体性能展示如下：

1. Yulun Zhang, Kunpeng Li, Kai Li, Bineng Zhong, and YunFu. Residual nonlocal attention networks for image restora-tion.ICLR, 2019.
2. Jacob Ziv. On universal quantization.IEEE Transactions onInformation Theory, 1985.3,5,6
3. Aaron Van Den Oord, Nal Kalchbrenner, Oriol Vinyals,Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. Con-ditional image generation with pixelcnn decoders.NIPS,2016.3,5,6