A Unified End-to-End Framework for Efficient [论文笔记]

论文链接

原文地址: A Unified End-to-End Framework for Efficient

论文代码: https://github.com/liujiaheng/compression

Abstract

作者提出的高效深度图像压缩的统一框架(EDIC)

• 通道注意模块 : 充分利用潜在表示之间的通道关系.
• 高斯混合模型 : 熵模型引入了高斯混合模型, 提高了比特率估计的准确性.
• 解码器端增强模块 : 进一步提高图像压缩性能.

且高效深度图像压缩框架(EDIC)很容易与深度视频压缩框架(DVC)相结合.

Proposed Methods

A. 图像压缩的整体架构

作者提出了高效深度图像压缩框架 (efficient deep image compression framework)，称为EDIC. 该网络有4个模块, 即编码器网络 (encoder), 解码器网络 (decoder), 超编码器网络 (hyper-encoder), 超解码器网络 (hyper-decoder).

编码器网络以原始图像 $x$ 为输入, 生成图像的潜在表示 $y$. 随后潜在表示 $y$ 被量化为 $\hat{y}$ . 在算术编码后, 量化的潜在表示 $\hat{y}$ 被发送到解码器网络, 以获得重建图像 $\hat{x}$ .

图像压缩方法, 旨在给定比特率的目标下获得高质量的重建图像，并且通过使用熵模型来估计比特率, 因此建立准确的熵模型至关重要. 文中应用超编码器和超解码器模块来估计熵模型的参数.

具体而言, 超编码模块基于潜在表示 $y$ 获得超先验信息并将其编码为超潜在表示 $z$ . 同样的, 超潜在表示 $z$ 被量化为 $\hat{z}$ . 在经过算术编码后, 超解码器将超潜在表示 $z$ 作为输入来构建超先验信息, 并估计熵模型的参数 $\Phi$ .

整体优化的损失函数为 :

$$L=\lambda D+R=\lambda d(x,\hat{x})+H(\hat{y})+H(\hat{z})$$

其中 $D$ 代表误差, $R$ 代表比特率. $\lambda$ 是权衡系数, $d(\cdot )$ 是误差度量(MS-SSIM). $H$ 代表潜在表示的 $\hat{y}$ 和 $\hat{z}$ 的比特率. 使用潜在表示的熵来近似表达比特率. 其中 $p_{\hat{y}|\hat{z}}(\hat{y}|\hat{z})$ 和 $p_{\hat{z}}(\hat{z})$ 分别表示 $\hat{y}$ 和 $\hat{z}$ 的分布.

$$\begin{gathered} H(\hat{y})=E[-lo{g}_2(p_{\hat{y}|\hat{z}}(\hat{y}|\hat{z}))] \\ H(\hat{z})=E[-lo{g}_2(p_{\hat{z}}(\hat{z}))] \end{gathered}$$

B. 通道注意力模块

作者提出了使用轻量的通道注意力技术来利用潜在表示 $\hat{y}$ 和 $\hat{z}$ 中的通道注意力. 具体结构如下:

其中输入特征图为 $X,X\in {\mathbb{R}}^{I\times J\times C}$ , $I,J,C$ 分别代表特征图的高,宽和通道数.

首先, 使用全局平均池化(GAP)来获得 channel-wise 统计 $t\in {\mathbb{R}}^C$ , 具体公式为:

$$t_c=\frac{1}{I\times J}\sum _{i=1}^I\sum _{j=1}^J{x}_c(i,j)$$

其中 $t_c$ 代表 t中的第c个元素. $x_c(i,j)$ 代表输入特征图 $X$ 第 $c$ 个通道的特定位置 $(i,j)$ 的值. 然后运用几个非线性变换来捕获 channel-wise 关系. 每个非线性变换层可描述为 :

$$s=\sigma (W_2)\delta (W_{1}t)$$

其中 $s$ 指输出 channel-wise 的注意力值，$W_1={\mathbb{R}}^{\frac{C}{r}\times C}$ 和 $W_2={\mathbb{R}}^{C\times \frac{C}{r}}$ 表示全连接层(FC)，$\delta$ 是非线性变换的ReLU激活函数，$\sigma$ 表示sigmoid函数激活。为了降低维数，将r设置为16。此外，在实现中还添加了残差运算。

所提出的通道注意力模块集成在编码器和超编码器网络中, 可利用通道关系进行高质量压缩.

C. 高斯混合模型的熵估计

在基于学习的图像压缩方法中，准确的比特率估计至关重要. 之前的许多方法将潜在表示 $\hat{y}$ 建模为高斯分布:
$$p_{\hat{y}|\hat{z}}(\hat{y}|\hat{z})\sim \mathcal{N}(\mu ,\sigma )$$

但高斯模型的对于复杂内容的表达能力有限. 因此利用高斯混合模型来进一步提高图像压缩系统的效率, 具体公式为:
$$p_{\hat{y}|\hat{z}}(\hat{y}|\hat{z})\sim \sum _{i=1}^F{w}_i\mathcal{N}({\mu }_i,{\sigma }_i)$$

其中 $w_i$ 代表不同高斯模型的权值, $F$ 为高斯模型的个数. 如图所示:

作者设计了三个卷积层和两个LeakyReLU层来估计高斯混合模型的参数 $\Phi$。在实现中 $F$ 被设为2。所以将GMM模块的输出通道数设为 5×N，前 4×N 个通道分别用于估计两个高斯模型的均值和方差. 在最后N个通道的输出上添加了一个sigmoid层, 为了估计每个高斯模型的权重. 如果一个高斯模型的权值为 $w_i$，那么另一个高斯模型的权值为 $1-w_i$ .

下图分别显示了 原始图像, 单高斯模型/混合高斯模型的潜在表示 $\hat{y}$ 的比特分配图, 两者间的差异

D. 解码端增强网络

由于提出的压缩方案是一个有损过程, 重建图像不可避免地有压缩伪影. 为了进一步提高重构质量，在图像重构后的解码器端引入了增强模块. 本文引入残差块来学习图像的高频信息进行图像压缩.

首先添加一个卷积层，将信道维数从3增加到32。然后，将三个增强块应用到卷积层的输出. 每个增强块有三个残差块。最后，添加卷积层和残差运算来获得重建图像.

左边是经过解码端增强后的重建图像, 右侧是学习得到的残差图像, 可以看到学习到的残差图像主要包含高频信息, 这意味着解码端侧增强模块有助于预测高频成分.

将其运用到视频压缩框架, 其中 $\{x_1,x_2,...,x_{t-1},x_t\}$ 代表当前视频序列. $x_t$ 为 $t$ 时的帧, ${\hat{x}}_t$ 为重建帧. $m_t$ 和 $r_t$ 分别代表运动信息和残差信息. 比特率估计模块是用来估计潜在表示的比特率的. “Q”表示量化.整体框架如下所示:

1. 运动估计与运动压缩 : 利用CNN模型来预测光流，光流代表运动信息 $v_t$ . 我们不直接对运动信息 $v_t$ 进行编码，而是将 $v_t$ 发送到运动压缩模块的编码器网络得到 $m_t$ , 然后将 $m_t$ 量化, 利用运动压缩模块的解码器来重构运动信息 ${\hat{v}}_t$ .
2. 运动补偿, 残差压缩与帧重建 : 运动补偿模块以前一帧重构后的 ${\hat{x}}_t$ 和运动信息 ${\hat{v}}_t$ 为输入, 得到预测帧 ${\overline{x}}_t$, 并认为这一帧尽可能接近当前帧 $x_t$ . 之后, 我们使用原始帧 $x_t$ 和 ${\overline{x}}_t$ 来获取残差信息 $r_t$, 其中 $r_t=x_t-{\overline{x}}_t$ . 残差压缩模块的编码器对残差信息 $r_t$ 进行编码, 得到潜在表示 $y_t$ . 同样，残差压缩模块的解码器对残差信息 ${\hat{r}}_t$ 进行重构, 然后, 可以得到最终的重构帧 ${\hat{x}}_t$, 其中${\hat{x}}_t={\hat{r}}_t+{\overline{x}}_t$ .
3. 框架优化 : 通过最小化以下 Rate-Distortion 权衡来优化整体框架 $$L_t=\lambda D_t+R_t=\lambda d(x_t,{\hat{x}}_t)+H({\hat{r}}_t)+H({\hat{m}}_t)$$ 其中$L_t$为当前时间 t 时的损失, $d(\cdot ,\cdot )$ 为当前帧 $x_t$ 与重构帧 ${\hat{x}}_t$ 之间的损失, $H({\hat{r}}_t)$ 和 $H({\hat{m}}_t)$ 为残差信息的潜在表示 ${\hat{r}}_t$ 和运动信息的潜在表示 ${\hat{m}}_t$ 的比特率, 通过比特率估计模块来估计.

Results

在图像压缩方面, 采用Flick.com上的20745幅高质量图像. 随机取 256×256 裁剪好的 patch 进行训练 . 为了进行性能评估, 计算了速率失真(RD)性能, 该性能是对柯达图像数据集中的所有图像进行评估 .

图像压缩的表现和细节

当模型用质量指标如 MS-SSIM 损失函数进行优化时, 我们采用 $\lambda$ 为 8192 的 MSE 损失函数优化的模型作为我们的预训练模型 . 并用不同的λ值(即16、32、64、128、256、384、512)对预训练模型进行微调 . 以 $1\times 10^{-5}$ 的学习率训练模型, 迭代50万次.

当使用MS-SSIM度量时，所提出的EDIC方法和其他图像压缩方法的率失真曲线对比。