CAE-ADMM：通过基于ADMM的修剪在压缩自编码器中进行隐式比特率优化

论文

论文

代码

代码

概述

  作者使用交替方向乘子法（ADMM）来优化失真和压缩效率之间的权衡。
  因为传统方法都是使用熵估计器来评估比特率，从而在图像质量和压缩率之间找到平衡点。
  本文使用的ADMM（Alternative Direction Method of Multipliers）是一种优化算法。ADMM的目的就是把数据表示为稀疏向量，稀疏向量就是其中大部分元素都是零的向量，这样就可以压缩数据了。还可以自适应地调整稀疏向量中的非零系数数量，以获得最佳的压缩效果。
  本文的目标是最小化：
$$d(\mathbf{x},\hat{\mathbf{x}})+\beta \cdot R(\hat{\mathbf{z}})$$
  d是输入图像 x 和重建 x之间的失真 $\hat{x}$ ,β > 0是控制比例，R是测量量化后的潜在编码 $\hat{z}$的比特率。
  作者使用交替方向乘子方法模块替换了熵估计器，对潜在编码进行的激进修剪鼓励稀疏性，并因此有助于比特率的优化。实验表明，这种修剪范式本身能够在隐式优化熵率的同时，在SSIM和MS-SSIM两种度量标准下产生比原始CAE和其他传统编解码器更好的结果。
  作者提出的CAE—ADMM框架结构如下图：

  模型架构如图所示，是对Lucas Theis提出的CAE进行的修改。编码器和解码器由卷积层组成。输入图像首先通过三个块进行下采样，每个块包含一个卷积层、一个批量归一化层和一个PReLU层。在15个残差块后，应用两个额外的下采样卷积块，其中最后一个卷积块生成z。量化器Q然后对其进行量化，并输入解码器，其架构与编码器相似。
   其中conv_same是边界镜像填充为1，卷积核3×3，步长为1的卷积核。conv_down边界镜像填充为2，卷积核5×5，步长为2的卷积核。sub_pixel里面有一个卷积核尺寸为1×1，步长为1的卷积核，PixelShuffle是像素重组，上采样，类似反卷积，可以放大图像尺寸。Boottleneck残差块，基于resnet的原理。

方法

  首先是典型的CAE（编码器 E、解码器 D 和量化器 Q 组成）结构：$$\begin{array}{rlr}E:{\mathbb{R}}^n\to {\mathbb{R}}^m,& & \\ D:{\mathbb{R}}^m\to {\mathbb{R}}^n,& & \\ Q:{\mathbb{R}}^m\to {\mathbb{Z}}^m.& & \\ & & \end{array}$$
  编码器E将原始图像x∈${\mathbb{R}}^n$映射到“潜在表示” （经过编码器处理后的特征向量一般具有更高的抽象性和更好的表达能力，能够更好地反映输入数据的本质特征）z=E(x)。量化器Q将z的每个元素映射到$\mathbb{Z}$，从而产生图像的压缩表示$\hat{z}$=Q(z)。最后，解码器D尝试从$\hat{z}$中的信息重构原始图像$\hat{x}$=D($\hat{z}$)。
  本文的目标可以由下式表达：
$$\underset{E,D}{\mathrm{argmin}}d(\mathbf{x},D\circ Q\circ E(\mathbf{x}))+\beta \cdot R\circ Q\circ E(\mathbf{x})$$
  总体来说就是最小化原始图像x和重构图像$\hat{x}$之间的误差，以及最小化比特率R。

量化器

  量化器Q采用了Lucas Theis等人提出的设计，其灵感来自于Toderici等人提出的随机二进制版本。
Q ( t ) = ⌊ t ⌋ + ϵ , ϵ ∈ { 0 , 1 } , Q(t)=\lfloor t\rfloor+\epsilon,\epsilon\in\{0,1\},_{} Q(t)=⌊t⌋+ϵ,ϵ∈{0,1},​

  方括号表示向下取整运算符， ϵ=1时满足 P ( ϵ = 1 ) = t − ⌊ t ⌋。这公式的意思是，如果t=1.7，那么t − ⌊ t ⌋的值是0.7，也就是ϵ = 1的概率是0.7，Q(t)=2。如果t=1.3,那么t − ⌊ t ⌋的值是0.3，也就是ϵ = 1的概率是0.3，大概率ϵ = 0。Q(t)=1。也就是说t小数部分大于0.5的编码为大于t的最小整数，t小数部分小于0.5的编码为小于t的最大整数。
  也就是说：
$$Q(t)=\{\begin{array}{ll}\lfloor t\rfloor +1& t-\lfloor t\rfloor >0.5\\ \lfloor t\rfloor & t-\lfloor t\rfloor <0.5\end{array}$$
  但是本文章里面的模型结构使用了多次的归一化Batch Normalization，以及relu，所以t中的数值都在0-1之间的。所以上面式子可以总结为：
$$Q(t)=\{\begin{array}{cc}1& t>0.5\\ 0& t<0.5\end{array}t\in [0,1]$$

  为了使量化器可微分，我们用它的期望定义它的梯度：
$$\frac{\partial }{\partial t}Q(t)=\frac{\partial }{\partial t}\mathbb{E}[Q(t)]=\frac{\partial }{\partial t}t=1$$

优化问题的解决方案

  作者提供一种替代方法来优化 R 而无需使用熵估计器H。我们可以通过以下方式重新表述 R:
$$\begin{array}{rlr}& & \\ & R(\hat{\mathbf{z}}& )=\mathrm{card}(\hat{\mathbf{z}})=\mathrm{card}(Q\circ E(\mathbf{x})),\\ & & \end{array}$$
  card(·) 计算非零元素的数量,希望编码器生成的 z 的非零元素数量少于所需数量$\ell$，我们可以将问题改写为 ADMM 可解决的问题:
$$\begin{gathered} {\mathrm{argmin}}_{E,D}d(\mathbf{x},\hat{\mathbf{x}})+g(\mathbf{Z}), \\ s.t.Q\circ E(\mathbf{x})-\mathbf{Z}=0. \end{gathered}$$
  其中指标函数 g(·) 定义为:
$$g(\mathbf{Z})=\{\begin{array}{ll}0& \text{if}\mathrm{card}(\mathbf{Z})\le \ell ,\\ +\infty & \text{otherwise.}\end{array}$$
  注意，$\mathbf{U}$和$\mathbf{Z}$都被初始化为全零，而$\mathbf{Z}$是 S = { Z ∣ card ⁡ ( Z ) ≤ ℓ } \mathbf{S}=\{\mathbf{Z}|\operatorname{card}(\mathbf{Z})\leq\ell\} S={Z∣card(Z)≤ℓ}的元素。通过引入对偶变量 $\mathbf{U}$ 和正则化因子 ρ > 0，我们可以将上述问题拆分为两个子问题。第一个子问题是：
$${\mathrm{argmin}}_{E,D}d(\mathbf{x},\hat{\mathbf{x}})+\frac{\rho }{2}\Vert Q\circ E(\mathbf{x})-{\mathbf{Z}}^k+{\mathbf{U}}^k{\Vert }_F^2,$$
其中k为当前迭代次数,$\Vert \cdot {\Vert }_F^2$是Frobenius范数，这是神经网络带有L2正则化的损失函数，可以通过反向传播和梯度下降进行求解。第二个子问题是：
$${\mathrm{argmin}}_{Z}g(\mathbf{Z})+\frac{\rho }{2}\Vert Q^{k+1}\circ E^{k+1}(\mathbf{x})-\mathbf{Z}+{\mathbf{U}}^k{\Vert }_F^2.$$
这个问题的解决方案是由 Boyd 等人在 2011 年推导出来的：$${\mathbf{Z}}^{k+1}={\Pi }_{\mathbf{S}}(Q^{k+1}\circ E^{k+1}(\mathbf{x})+{\mathbf{U}}^k),$$
其中，ΠS(·)表示欧几里德投影到集合S上。一般来说，对非凸集合进行欧几里德投影是困难的，但是Boyd等人证明了最优解是保留$Q^{k+1}\circ E^{k+1}(\mathbf{x})+{\mathbf{U}}^k$中$\ell$个最大元素，并将其余元素设为零。最后，我们将按照以下策略更新对偶变量U：$${\mathbf{U}}^{k+1}={\mathbf{U}}^k+Q^{k+1}\circ E^{k+1}(\mathbf{x})-{\mathbf{Z}}^{k+1}.$$这三个步骤共同构成了 ADMM 剪枝方法的一次迭代。
  下面是完整的步骤：

结果

  使用批量大小设置为 32 的 Adam 优化器来解决第一个子问题。学习率设置为 4 · $10^{-3}$，每次损失在十个时期内没有下降时减半。每 20 个 epoch，应用 ADMM 剪枝方法的第二步和第三步。第二步保留的元素个数比例设置为10%。
  Bpp（比特率）和ˆz中零元素的比例，以及修剪前后ˆz中的元素总数。

代码部分与修改

  num_resblocks后面缺少参数。我们传入16

  ADMM算法中的Z,U我们改成16，与量化后的矩阵尺寸相同。

  下面同样都给改了


  我发现batchsize太大了会报爆显存的错误，可以把batchsize调小一点，例如2。但是训练完一轮之后还会产生爆显存的错误，经过调试之后发现是在验证阶段报的错误。
  原代码：

  验证阶段batch_size 这部分 batch_size=testset.__len__()直接把所有图片放进去了，原作者的显卡的显存足够大。但是我的显存不够所以要修改一下。把batch_size也改成2。
  一旦把batch_size改成2之后就又会出现一个问题，那就是保存图片的时候。
  下面是train验证部分的代码：

  下面utils部分的代码：

  我们可以看到验证部分，是使用了kodak数据集的24张图片。但是本网络结构是压缩128×128分辨率的图片，kodak数据集的图片大小为768×512，所以我们需要对其裁剪，垂直裁剪6块，水平裁剪4块。一张768×512图片裁减为24张128×128尺寸的图片。
  下面保存图片的这一部分，数字不能大于batch_size，这里我就选择只保存第一张图片。

结果图

下面是结果图：
  下面是我使用2286张图片训练200轮后的效果图。

训练完后的模型链接：https://pan.baidu.com/s/16WDL_JJjxVM4W9pc5I_Bbw?pwd=a2u9
提取码：a2u9