交替方向乘子法的自编码器——基于压缩自编码的ADMM修剪的隐式比特率优化
2019年
Haimeng Zhao
摘要:
本文介绍了一种采用交替方向乘子法(ADMM,Alternative Direction Method of Multipliers)的ADMM剪枝压缩自动编码器。具体来说,我们的方法中的ADMM是为了提高稀疏性来隐式优化比特率,不同于以往研究中使用的熵估计器。在公共数据集上的实验表明,该方法在SSIM/MS-SSIM指标上优于原CAE和部分传统编解码器,推理速度合理。
介绍:
对于目标函数不可微的问题,Theis[1]提出的一种熵编码和熵率估计的结合的方法,其中训练了一个参数化熵估计器与编码器/解码器。随着取上限过程和高斯尺度混合(GSMs)的运用,整个过程是可微分的。
作为上述方法的一种替代方法,我们用乘子模块的交替方向法代替熵估计器,其中对潜在编码的渐进剪枝促进了稀疏性,因此有助于比特率的优化。实验表明,在SSIM和MS-SSIM两种情况下,该剪枝范式本身能够隐式地优化熵率,同时与原始CAE和其他传统编解码器相比,具有更好的效果。
相关工作:
熵估计模型包括:高斯尺度混合模型、生成模型、上下文模型。对于非自编码方法的工作,人们对GAN[9]、GDN(广义分裂归一化)[5]和RNN[10]的使用也越来越感兴趣。传统的凸优化ADMM算法[11]首先通过Han等人提出的方法在神经结构搜索(NAS)领域得到了应用。
提出的方法:
与Theis的方法[1]类似,CAE-ADMM使用卷积层作为我们的编码器和解码器的基础。解码器镜像编码器的结构以保持对称性,不同之处在于未使用Shi等人提出的亚像素卷积层来执行上采样。
量化:随机二值化(简单且计算高效)(Toderici等)
ε决定是输出为输入的ground还是ceiling,且ε= 1 的概率为:
量化函数是不可微的,在后向传播中,利用期望来求导数:
优化问题的求解(就是先用一个式子来代替rate,再介绍ADMM修剪算法来求解这个率失真优化损失函数):
不用熵H的方法,我们重新制定R(rate)
其中card(·)计算非零元素的个数。如果我们想让编码器生成的z的非零元素的数量比期望的数字l少,我们可以将问题重新表述为一个admm可解的问题[16]:
其中指示器函数g(·)定义为:
注意U和Z都初始化为全零,Z是S中的一个元素,S为:
通过引入对偶变量U和惩罚因子ρ> 0,我们可以把上面的问题分成两个子问题。第一个子问题是:
其中k为当前迭代次数。这是L2正则化神经网络的损失,可以通过反向传播和梯度下降来解决。第二个子问题是:
这个问题的解决方案是由Boyd等人推导出来的(2011 [18]):
其中Πs(·)代表了集合S上的欧几里得投影。一般来说,非凸集上的欧几里德投影是困难的,但是Boyd等人[18]证明了最优解是保持括号中的最大的l个元素,其余的设为零。最后,我们将使用以下策略更新对偶变量U:
这三个步骤一起构成了ADMM剪枝方法的一个迭代。算法1显示了完整的步骤:
每次迭代把QE(x)中最大的l个元素保留下来,其余的设置为0,枝剪体现在这里。
实验:
模型框架:
训练:
我们使用Adam优化器[19]将batch大小设置为32来解决第一个子问题。学习率设置为4·10exp(-3),10个epoch不下降的话,学习率减半。每隔20个epoch,应用ADMM剪枝方法的第二步和第三步。作为反向传播一部分的距离函数是MSE与PSNR/SSIM/MS-SSIM的可微版本的线性组合,训练首先由一个缩放的MSE单独进行。步骤2中要保留的元素数量的比例设置为10%。为了支持bpp的细粒度调优,我们修改了编码器的最后一层。所有过程都在PyTorch中实现且开放源码。每个模型在4个NVIDIA GeForce GTX 1080Ti gpu上训练300个epoch。
数据集和预处理:
我们使用BSDS500[20]作为训练集,其中包含500幅481×321自然图像。图像随机裁剪到128×128,水平和垂直翻转,然后归一化。对于测试集,我们使用Kodak PhotoCD数据集,其中包含24张768×512的图像。
结果和讨论:
图2.比较不同方法对SSIM和MS-SSIM在Kodak PhotoCD 数据集上的应用。注意Toderici等人[21]用RNN结构代替熵编码,而CAE-ADMM (our)用剪枝方法代替熵编码。
CAE是Theis那一篇,利用了熵编码的方法。
在消融研究中,我们将相同的训练程序应用于相同的模型,一个有修剪计划,另一个没有,来验证ADMM模型的有效性。然后,我们计算了混合数据集(Urban100[22]、Manga109[22]和Kodak PhotoCD中的768*512帧图像)中的平均bpp和零元素比率。结果可以在表1中看到,图4中可以看到更直观的示例图像。
表1. Bpp &剪枝前后z^中0个元素的比例。对于这两种统计,95%置信区间是建立在样本大小为233(混合数据集的大小)的情况下的。
图4. kodim21剪枝前后潜在编码的比较。为了清晰起见,我们在归一化之前将feature map中的零值标记为黑色。
在推理速度方面,从表2可以看出,我们的CAE-ADMM与传统的编解码器相比具有可接受的推理速度,同时保持了SSIM的优良质量。
表2. 混合数据集上SSIM的95%置信区间和CAE-ADMM模型的推理速度(推理时间为所有128x128的patch的总和,批量大小= 8),在2个gpu上与传统的编解码器相比较。
结论:
本文提出了一种基于admm剪枝的压缩自编码器(CAE-ADMM)[23],它是一种基于深度学习的有损图像压缩熵估计技术的替代方案。在多个数据集上的测试显示了比原始CAE模型和其他当代方法更好的结果。通过研究学习的潜在编码细节,我们进一步探讨了基于admm的剪枝方法的有效性。
进一步的研究可以集中在开发一个更有效的剪枝方法,例如,从强化学习领域引入思想。此外,E、D和Q的结构可以进一步优化,以提高速度和精度。
Learning Content-Weighted Deep Image Compression
MuLi
2019
摘要:
基于学习的有损图像压缩通常涉及率失真性能的联合优化。现有的方法大多采用空间不变的比特长度分配,并结合离散熵近似来约束压缩率。然而,信息内容在空间上是变化的,其中结构复杂和突出的区域通常对图像压缩更为重要。考虑到图像内容的空间变化,提出了一种基于内容加权的编解码器模型,该模型包含一个重要性图子网络来生成局部自适应比特率分配的重要掩模。因此,重要掩模的求和可以作为压缩率控制的熵估计方法。此外,潜在编码的量化表示和重要性图仍然是空间相关的,可以使用算术编码对其进行无损压缩。为了有效地压缩编码,我们提出了一个trimmed(裁剪)卷积网络来预测量化码的条件概率。实验表明,与传统的深度有损图像压缩方法相比,该方法具有良好的视觉效果。
介绍:
Balle等人的[5]只是采用基于上下文的自适应二进制算术编码框架(CABAC)[8]进行熵编码。考虑到维护条件概率表的实际可行性,CABAC只使用两个最近的码作为上下文。尽管如此,在熵编码中加入更大的上下文可以获得更好的压缩性能。
为了利用局部空间上下文,我们采用算术编码框架,提出了一种卷积熵预测模型,用于量化码和重要图的上下文预测当前符号。
我们提出了一种用于算术编码的trimmed卷积网络(TCAE),其中卷积内核经过特殊的修剪,以适应压缩顺序和上下文依赖关系。然后,通过全卷积网络,可以有效地对所有符号进行概率预测。通过叠加多个裁剪后的卷积层,TCAE可以在保持计算效率的同时对大上下文进行建模。在此基础上,提出了一种倾斜的TCAE模型,将三维编码图中的编码划分为多个斜面。然后可以安全地对每个斜面内的符号进行并行解码,从而大大加快解码过程。
这篇论文是我们先前工作[12]的一个实质性的扩展。与[12]相比,在编码器和解码器中引入了一种结合dense block的改进的网络结构。用一种基于学习信道的多值量化方法代替二进制量化方法,实现了编码器特征的自适应判断。此外,为了更好地学习重要图子网,采用了两阶段松弛算法。最后,我们引入了一个TCAE和一个倾斜的TCAE来有效地对算术编码中的大上下文进行建模。这项工作的贡献摘要如下:
- 提出了一种用于有损图像压缩的内容加权编解码器模型。这里,一个学习的逐通道量化进行编码器特性的判断,引进重要性地图子网络指导局部自适应比特分配,生成的重要性掩模的总和作为压缩率的估计,且提出了一种学习重要图子网的两阶段松弛算法。
- 针对算术编码中的大上下文建模问题,提出了一种TCAE网络。利用裁剪卷积,通过全卷积网络可以有效地预测量化码的条件概率。并进一步引入倾斜TCAE模型来加速解码过程。
- 实验表明,该方法在抑制视觉伪影的同时,有效地恢复了图像的显著结构和丰富的细节。此外,与现有的图像编码标准[1]、[2]、[3]和深度模型[5]、[9]、[10]、[11]相比,我们的方法性能良好。