CAE-P:基于ADMM剪枝的图像压缩自编码器

赵海萌

前言

图像压缩自编码器CAE，亦即编码器E及解码器D，其训练过程可以转化为一个优化问题，即对图像失真率及图像编码比特数的最小化。有损图像压缩面临着失真率及压缩率的权衡问题，因此可以将上述优化问题表述为：

minE，Dd+βR

其中d表示重构图像与原图像之间的差距，R表示图像编码比特数，而β>0则控制上述两个因素之间的平衡。解决这个优化问题的过程中会遇到许多困难，其中最为重要的一个是如何表征图像编码比特数R。因此，本项目针对R的优化，提出利用ADMM（AlternatingDirectionMethodofMultipliers）算法对CAE的表示层进行剪枝，即直接减小R，避开了额外训练信息熵估计器的麻烦，并遵循训练、剪枝、重训练的顺序，迭代地对CAE进行训练（优化d）和剪枝（优化R），直至达到目标要求。

本文提出了CAE-P（CompressiveAutoEncoderwithPruning）模型，相较于现有的CAE模型，CAE-P模型显得更为简单直接，更易实现且参数量更小。实验中，CAE-P模型在MS-SSIM（Multi-scaleStructuralSimilarityIndex）、SSIM等指标下均超越了现有的图像压缩算法。

CAE-P模型

一个基本的图像压缩自编码器CAE由3部分组成：编码器E、解码器D和量化器Q。图1为CAE-P的详细模型架构，“Convk/spP”表示卷积核大小为k×k、步长为s，并采用大小为P的镜像Padding的卷积层，“ConvDown”表示将宽和高减半的卷积层。

在对编码器E、解码器D、量化器Q进行选择及构造，并解决优化问题后，我们便得到了基于ADMM剪枝的图像压缩自编码器CAE-P，它的基本结构及运作方式如图2所示。

原始图像x经过由卷积残差块构成的编码器E编码，转化为一组潜在表示形式z。在ADMM训练阶段，黑色和绿色通路激活，特征图z经量化器Q量化后得到，输入同样由卷积残差块构成的解码器D，解码器D从中重构出图像，ADMM算法迭代地最小化重构失真率d（x，），并对z剪枝，迫使其稀疏化。ADMM训练阶段结束后，绿色通路关闭，红色通路开启，进入重训练阶段。此时由于ADMM算法并不能保证z中仅剩余小于等于ℓ个非零元素，只能使其他元素接近于零，因此需要对z中的近零元素强制剪除，再经量化器Q量化后，得到最终编码形式，并将此最终编码形式输入解码器D中，得到重构图像。为了减小强制剪除对重构图像的影响，此时的重训练便只需优化minE，Dd（x，）。遵循训练、剪枝、重训练的步骤反复训练，可以得到同时具有高图像重构质量及高压缩率的图像编码。这便是本文所提出的CAE-P模型的基本架构及训练方法。

图2为CAE-P模型的训练策略图。其中，Encoder和Decoder均采用基于CNN的残差块构建，Code为最终的压缩后编码。ADMM训练过程中，黑色和绿色通路激活，ADMM算法在最小化重构失真率的同时对z剪枝，迫使其稀疏化。ADMM训练结束后，黑色和红色通路激活，对z中的近零元素强制剪除，开始重训练。

实验过程

模型架构

CAE-P模型由基于CNN和残差块的编码器、解码器组成。原始图像首先经过3次下采样，每次包含下采样卷积、批归一化BN层和PReLU激活函数;接下来经过15层Bottleneck残差块，其中每次卷积后均连接BN层和PReLU;最后再经过2层下采样卷积，得到z。

接下来经过量化器得到，量化器的梯度如下，设置为1。

量化后的编码进入解码器，解码器的架构与编码器基本镜像。编码首先通过2层上卷积Sub-Pix层，每层包含卷积、BN层、PixelShuffle和PReLU激活函数;然后经过15层Bottleneck残差块，内部结构与编码器中的相同;最后再经过2层上卷积Sub-Pix层，最后通过Tanh激活函数，使激活值限制在（-1，1）之间，最后将其线性映射至（0，1）之间。

目标函数d的选择

目前，常用的d的选择包括MSE、MS-SSIM等。我们尝试了多种设计后，最终选择d（x，）=100·（1-MSSSIM（x，））+100·（1-SSIM（x，））+（45-PSNR（x，））+MSE（x，）。所有模型均以此为目标函数d训练。

模型训练

研究采用Adam作为优化器，BatchSize设置为32，对ADMM优化问题的第一步进行优化。初始学习率设置为4·10-3，并随着训练过程动态衰减：每出现10个Epoch平均损失函数不下降，学习率减半。WeightDecay设置为一个较小值1·10-10。每经过20个Epoch，执行1次ADMM算法的第2步和第3步，第2步中的保留元素比例设置为10%。对于不同bpp的模型，修改编码器的最后一层及解码器的第1層的结构，进行fine-tune。所有代码在PyTorch框架下实现，每个模型在4块NVIDIAGeForceGTX1080TiGPU上并行化训练了300个Epoch，并在GitHub上开源。

数据集及其预处理

实验选择BSDS500作为训练数据集，包含500张481×321的图片。每次输入网络时图像首先被随机裁剪出一块128×128大小的Crop，随机进行水平和竖直翻转，最后归一化至（0，1）。对于测试集，实验采用图像压缩领域内通用的测试数据集KodakPhotoCD2，包含24张768×512的照片。

实验结论

研究提出了CAE-P模型，避开了额外训练熵编码器的繁琐与优化的不直接，引入ADMM算法对编码层剪枝，优化更为直接，更易实现且参数量更小。在现有的使用熵编码器的方法之外开辟并探索了一条新的道路：利用剪枝方法。实验表明，实验模型在多项指标下均超越了现有的图像压缩算法。

研究意义

本研究致力于改进图像有损压缩算法的压缩率与图像重构质量，首次将最先进的神经网络架构搜索技术迁移于其中，实现了超越现有图像压缩算法的性能，为图像更高效地存储与传输提供了基础。

创新点

本研究提出了CAE-P模型，首次在图像压缩领域中引入剪枝方法，在使用熵估计器的方法之外开辟并探索了应用架构搜索方法（剪枝）的可能性，为未来引入更多更精巧的架构搜索方法作了铺垫。实验表明，该模型在多项指标下均超越了现有的图像压缩算法，尤其是研究改进对象CAE。