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

• 论文链接
• Abstract
• Proposed Methods
• • A. 图像压缩的整体架构
  • B. 通道注意力模块
  • C. 高斯混合模型的熵估计
  • D. 解码端增强网络
• Results
• • 图像压缩的表现和细节

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

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

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

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

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

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

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

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

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

整体优化的损失函数为 :

L = λ D + R = λ d ( x , x ^ ) + H ( y ^ ) + H ( z ^ ) L = lambda D + R = lambda d(x, hat x) + H(hat y) + H(hat z) L=λD+R=λd(x,x^)+H(y^​)+H(z^)

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

H ( y ^ ) = E [ − l o g 2 ( p y ^ ∣ z ^ ( y ^ ∣ z ^ ) ) ] H ( z ^ ) = E [ − l o g 2 ( p z ^ ( z ^ ) ) ] H(hat y) = E[-log_2(p_{hat y|hat z}(hat y|hat z))] \ H(hat z) = E[-log_2(p_{hat z}(hat z))] H(y^​)=E[−log2​(py^​∣z^​(y^​∣z^))]H(z^)=E[−log2​(pz^​(z^))]

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

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

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

t c = 1 I × J ∑ i = 1 I ∑ j = 1 J x c ( i , j ) t_c = frac{1}{I times J} sum_{i=1}^{I} sum_{j=1}^{J} x_c(i, j) tc​=I×J1​i=1∑I​j=1∑J​xc​(i,j)

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

s = σ ( W 2 ) δ ( W 1 t ) s = sigma(W_2)delta(W_1 t) s=σ(W2​)δ(W1​t)

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

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

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

但高斯模型的对于复杂内容的表达能力有限. 因此利用高斯混合模型来进一步提高图像压缩系统的效率, 具体公式为:
p y ^ ∣ z ^ ( y ^ ∣ z ^ ) ∼ ∑ i = 1 F w i N ( μ i , σ i ) p_{hat y|hat z}(hat y|hat z) sim sum_{i=1}^{F}w_imathcal{N}(mu_i,sigma_i) py^​∣z^​(y^​∣z^)∼i=1∑F​wi​N(μi​,σi​)

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

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

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

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

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

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

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

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

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

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

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