图像压缩-《Learned Image Compression with Discretized Gaussian Mixture Likelihoods and Attention Modules》

论文地址：https://arxiv.org/abs/2001.01568
代码地址：https://github.com/ZhengxueCheng/Learned-Image-Compression-with-GMM-and-Attention

    基于论文《Joint Autoregressive and Hierarchical Priors for Learned Image Compression》的基础上进行了码率模型的优化。

一. 简介

• 从基线论文《Joint Autoregressive and Hierarchical Priors for Learned Image Compression》中可视化了压缩特征，并且发现特征图中依旧存在未捕获的结构冗余特征。因此，我们建议使用离散的高斯混合似然来对分布进行参数化，从而消除了剩余的冗余以实现准确的熵模型，从而直接导致所需的编码位数更少。
• 此外，网络架构中采用了注意力模块的简化版本，注意模块可以使学习的模型更加关注复杂区域，从而以中等训练复杂度提高我们的编码性能。

二. 内容

具体基础图像压缩基础内容可参考其他的博文，文章的主要创新描述的内容如下：

• a) baseline 框架具体来自于Balle论文《End-to-end optimized image compression》，内容可参考博文【1】
• b) Hyperior 框架同样来自于google实验室Balle出品《Variational Image Compression With A Scale Hyperprior》，内容参考博文【2】
• c) Joint 是基于Hyperior之上优化的，由Minnen发表的《Joint Autoregressive and Hierarchical Priors for Learned Image Compression》，具体内容可参考博文【3】
• d) 即本文提出的在 c) 的基础上，优化了对于潜在的特征点的建模模型，从单高斯分布的形式优化为混合高斯分布的形式。

整体的模型框架以及卷积核参数：

在整体的网络中，对比c）中的模型，该整体框架将1个5x5的卷积核分解成2个3x3的卷积核，根据图鸭科技的论文中描述，这种卷积核的变换可以达到相同的感受野，并且可以有一定的性能增益，并且在论文中，增加了简易注意力机制模块，通过1x1的卷积核增强通道间的信息交流。并且通过sigmoid函数生成重要性掩膜，使得网络更注重复杂纹理区域的

柯达数据集中的kodim21为例，对具有最高熵的通道的不同熵模型进行可视化。它显示了该熵方法提供了更灵活的参数化分布模型，具有更小的比例参数和更好的空间冗余减少，这直接导致了更精确的熵模型和更少的码率消耗问题。

    通常，预测平均值 $\mu$ 接近$y$。复杂区域的 $\sigma$ 大，则需要更多的比特进行编码。取而代之的是，平滑区域的$\sigma$较小，导致用于编码的位更少。 分层先验模型仍然在简单区域（例如kodim21的天空）中保留了空间冗余。同样，我们可视化联合熵模型。与分层先验相比，Joint通过添加通过5×5掩码卷积实现的自回归模型来捕获更多与归一化值的结构，以捕获与相邻元素的相关性。

    但是，联合熵模型并不是完美的，因为在上图的第二行（第四列）中仍然观察到一些空间冗余。尽管相邻元素已经用作上下文模型的输入，但是参数化分布无法表示充分利用相邻元素的上下文和信息以及附加位 $\hat{z}$。可能受到单一高斯分布的固定形状的限制。这促使考虑使用更灵活的参数化模型来实现任意可能性。因此，我们提出了高斯混合模型，即$$p_{\hat{y}|\hat{z}}(\hat{y}|\hat{z})=\sum _{k=1}^K{w}_i^{(k)}N({\mu }^{(k)},{\sigma }^{2(k)})$$通过 $k$ 调节高斯混合模型中，高斯分布的个数问题，其核心公式如下$$\begin{gathered} p_{\tilde{y}|\tilde{z}}(\tilde{y}|\tilde{z},{\theta }_h)=\sum _{k=1}^K{w}_i^{(k)}N({\mu }_i^{(k)},{\sigma }_i^{2(k)})\ast U(-1/2,1/2))(\widetilde{y_i}) \\ =c(\widehat{y_i}+1/2)-c(\widehat{y_i}-1/2) \end{gathered}$$特征图中指定位置，${\tilde{y}}_i$ 表示 $y$ 的第 $i$ 个元素，而${\mu }_i$表示 $\mu$ 的第 $i$ 个元素。 $k$ 表示混合高斯的数量。每个高斯模型都是具有3个参数的高斯分布，即权重$w_i^{(k)}$，均值${\mu }_i^{(k)}$和每个元素yi的方差${\sigma }_i^{2(k)}$，这三个参数都是通过先验架构以及自回归模块学习得到的.$c()$是累积函数。$\hat{y}$ 的范围会自动获悉，并且不会提前知道。为了获得稳定的训练，我们将 $y$ 的范围限制为[-255，256]，因为根据经验y不会超过该范围。对于−255的边缘情况，将$c（\widetilde{y_i}-1/2）$替换为零，即$c(+\infty )=0$。对于256的边缘情况，将$c(\widetilde{y_i}+1/2)$ 替换为一个，即$c(-\infty )=1$它为训练提供了数值稳定的实现。
    通过概率密度参数生成表示数值出现概率的概率密度函数，通过对概率密度函数的单位仓内积分得到特征点的概率情况。通过这种形式达到了获取自适应的出现概率值的目的。通过得到的概率数值进行以此进行熵率估计可以微分的目的。

三. 核心代码：

训练高斯混合

def entropy_parameter(tensor, inputs, num_filters, training):
  """tensor: the output of hyper autoencoder (phi) to generate the mean and variance
     inputs: the variable needs to be encoded. (y)
  """
  with tf.variable_scope("entropy_parameter", reuse=tf.AUTO_REUSE):

    half = tf.constant(.5)

    if training:
      noise = tf.random_uniform(tf.shape(inputs), -half, half)
      values = tf.add_n([inputs, noise])

    else: #inference
      #if inputs is not None: #compress
      values = tf.round(inputs)
        

    masked = masked_conv2d(values, num_filters*2, [5, 5], "A", scope='masked')
    tensor = tf.concat([masked, tensor], axis=3)
      
	""" 通过三个卷积层生成了对每个特征点进行建模所需要地建模参数，对于每个特征点需要三个高斯模型进行描述，每个高斯模型又需要三个参数，故而需要生成九个参数"""
    with tf.variable_scope("layer_0"):
      layer = tfc.SignalConv2D(
          640, (1, 1), corr=True, strides_down=1, padding="same_zeros",
          use_bias=True, activation=tf.nn.leaky_relu)
      tensor = layer(tensor)

    with tf.variable_scope("layer_1"):
      layer = tfc.SignalConv2D(
          640, (1, 1), corr=True, strides_down=1, padding="same_zeros",
          use_bias=True, activation=tf.nn.leaky_relu)
      tensor = layer(tensor)

    with tf.variable_scope("layer_2"):
      layer = tfc.SignalConv2D(
          num_filters*9, (1, 1), corr=True, strides_down=1, padding="same_zeros",
          use_bias=False, activation=None)
      tensor = layer(tensor)
    prob0, mean0, scale0, prob1, mean1, scale1, prob2, mean2, scale2 = \
             tf.split(tensor, num_or_size_splits=9, axis = 3)
    scale0 = tf.abs(scale0)
    scale1 = tf.abs(scale1)
    scale2 = tf.abs(scale2)
    probs = tf.stack([prob0, prob1, prob2], axis=-1)
    probs = tf.nn.softmax(probs, axis=-1)  //保证概率密度的总区间内积分的总和为1.
  
    # To merge them together
    means = tf.stack([mean0, mean1, mean2], axis=-1)
    variances = tf.stack([scale0, scale1, scale2], axis=-1)


	if training:
      """ 根据每个建模参数的生成三个高斯分布，并且根据边际分布获得待编码点在每个高斯分布中得到的概率情况，根据每个高斯概率密度函数的概率进行乘法计算，得到likelihoods"""
      # =======REVISION: Robust version ==========
      edge_min = probs[:,:,:,:,0]*dist_0.cdf(values + half) + \
                 probs[:,:,:,:,1]*dist_1.cdf(values + half) + \
                 probs[:,:,:,:,2]*dist_2.cdf(values + half)
      
      edge_max = probs[:,:,:,:,0]* (1.0 - dist_0.cdf(values - half)) + \
                 probs[:,:,:,:,1]* (1.0 - dist_1.cdf(values - half)) + \
                 probs[:,:,:,:,2]* (1.0 - dist_2.cdf(values - half))
      likelihoods = tf.where(values < -254.5, edge_min, tf.where(values > 255.5, edge_max, likelihoods))

      
      likelihood_lower_bound = tf.constant(1e-6)
      likelihood_upper_bound = tf.constant(1.0)
      likelihoods = tf.minimum(tf.maximum(likelihoods, likelihood_lower_bound), likelihood_upper_bound)

推理阶段高斯混合

"""通过获得三组的Mu,sigma,weight 参数推断得到总的pmf函数，概率质量函数，即离散情况下的概率密度函，通过accumulate函数叠加后得到CDF，之后输入算数编码器，得到编码结果。"""

mu = y_means_values[0, pad_size, pad_size, ch_idx, :] + minmax
sigma = y_variances_values[0, pad_size, pad_size, ch_idx, :]
weight = y_probs_values[0, pad_size, pad_size, ch_idx, :]

pmf = (0.5 * (1 + scipy.special.erf((samples + 0.5 - mu[0]) / ((sigma[0] + TINY) * 2 ** 0.5))) - \
                   0.5 * (1 + scipy.special.erf((samples - 0.5 - mu[0]) / ((sigma[0] + TINY) * 2 ** 0.5)))) * weight[0] + \
                  (0.5 * (1 + scipy.special.erf((samples + 0.5 - mu[1]) / ((sigma[1] + TINY) * 2 ** 0.5))) - \
                   0.5 * (1 + scipy.special.erf((samples - 0.5 - mu[1]) / ((sigma[1] + TINY) * 2 ** 0.5)))) * weight[1] +\
                  (0.5 * (1 + scipy.special.erf((samples + 0.5 - mu[2]) / ((sigma[2] + TINY) * 2 ** 0.5))) - \
                   0.5 * (1 + scipy.special.erf((samples - 0.5 - mu[2]) / ((sigma[2] + TINY) * 2 ** 0.5)))) * weight[2]
pmf_clip = np.clip(pmf, 1.0/65536, 1.0)
pmf_clip = np.round(pmf_clip / np.sum(pmf_clip) * 65536)
cdf = list(np.add.accumulate(pmf_clip))
cdf = [0] + [int(i) for i in cdf]
symbol = np.int(y_hat_value[0, h_idx, w_idx, ch_idx] + minmax )
encoder.encode([symbol], cdf)

四. 性能结果

在psnr性能指标上已经获得了和VVC相近的性能效果，MS-SSIM指标则传统的编码器效果好像比起端到端的压缩方式一直以来都差了一点，总的来说离散高斯混合模型在较低的计算cost的给整个编码框架带来了比较大的性能增益。