《END-TO-END OPTIMIZED IMAGE COMPRESSION》压缩模型

asdfghjuuuu

1847人浏览 · 2025-03-11 19:03:32

asdfghjuuuu · 2025-03-11 19:03:32 发布

一、输入

（一）预处理

读取png图像并将图像解码为RGB三通道格式

def read_png(filename):
  """Loads a PNG image file."""
  string = tf.io.read_file(filename)
  return tf.image.decode_image(string, channels=3)

RGB图像每个像素的取值范围为0-255。在进入分析变换前，需对输入数据进行归一化处理，将其映射到浮点数区间0，1。具体实现为将像素值除以255：

$x_{norm}=\frac{x}{255}$

代码：

self.add(tf.keras.layers.Lambda(lambda x: x / 255.))

（二）图像分块

在训练过程中，模型对预处理后的图像进行裁剪，生成固定大小256×256的图像块，裁剪后的图像块被送入分析变换网络。

代码：

def crop_image(image, patchsize):
    image = tf.image.random_crop(image, (patchsize, patchsize, 3))
    return tf.cast(image, tf.keras.mixed_precision.global_policy().compute_dtype)

二、分析变换

分析变换分由三层卷积网络构成，逐步降低空间分辨率并提取高阶特征。每一层包含以下操作：

1.卷积：使用特定大小的卷积核进行特征提取，生成特征图

2.下采样：对特征图进行下采样，通过步长控制空间分辨率缩减，有损压缩

3.GDN：引入非线性激活函数，对卷积和下采样后的特征图进行归一化处理。

（一）第一层

卷积：卷积核9×9，步长4

下采样：特征图尺寸降到原图 $\frac{1}{4}$

GDN：对第二层的输入进行归一化

（二）第二层

卷积：对第一层处理完的特征图再进行卷积，卷积核5×5，步长2

下采样：特征图进一步缩减，尺寸变为原图的 $\frac{1}{2}$

GDN：对第三层的输入进行归一化

（三）第三层

卷积：对第二层处理完的特征图再次进行卷积，卷积核5×5，步长2

下采样：再一次缩减特征图尺寸

无GDN

代码：

class AnalysisTransform(tf.keras.Sequential):
  """The analysis transform."""
  def __init__(self, num_filters):
    super().__init__(name="analysis")
    self.add(tf.keras.layers.Lambda(lambda x: x / 255.))
    self.add(tfc.SignalConv2D(
        num_filters, (9, 9), name="layer_0", corr=True, strides_down=4,
        padding="same_zeros", use_bias=True,
        activation=tfc.GDN(name="gdn_0")))
    self.add(tfc.SignalConv2D(
        num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
        padding="same_zeros", use_bias=True,
        activation=tfc.GDN(name="gdn_1")))
    self.add(tfc.SignalConv2D(
        num_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
        padding="same_zeros", use_bias=False,
        activation=None))

合成变换：分析变换逆过程，用于解压缩

代码：

class SynthesisTransform(tf.keras.Sequential):
  """The synthesis transform."""
  def __init__(self, num_filters):
    super().__init__(name="synthesis")
    self.add(tfc.SignalConv2D(
        num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
        padding="same_zeros", use_bias=True,
        activation=tfc.GDN(name="igdn_0", inverse=True)))
    self.add(tfc.SignalConv2D(
        num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
        padding="same_zeros", use_bias=True,
        activation=tfc.GDN(name="igdn_1", inverse=True)))
    self.add(tfc.SignalConv2D(
        3, (9, 9), name="layer_2", corr=False, strides_up=4,
        padding="same_zeros", use_bias=True,
        activation=None))
    self.add(tf.keras.layers.Lambda(lambda x: x * 255.))

三、量化

对特征图进行量化，量化为有损压缩，该模型采用均匀标量量化方法。

若用y表示需要被量化的数据，则量化过程为：

为了在训练过程中处理量化的不可微性问题，在训练时，用加性均匀噪声 $\Delta y$ 来模拟均匀量化

y_hat, bits = entropy_model(y, training=training)

在测试时使用真实的均匀量化

y = self.analysis_transform(x)
return self.entropy_model.compress(y), x_shape, y_shape

训练阶段量化近似

用均匀噪声U(-0.5,0.5)模拟量化误差

$y_{i}$ 是编码空间的元素，就是需要被编码的值， $\hat{y_{i}}$ 是四舍五入后的值, $\tilde{y_{i}}$ 是通过添加噪声后的值，

上述是三者的概率密度函数（PDF），离散的 $\hat{y_{i}}$ 是概率质量函数（PMF）

$\tilde{y_{i}}$ 的概率密度函数为

$p_{\tilde{y_{i}}}=\int_{-0.5}^{0.5}P_{q}\left ( y-u \right )du$

在所有整数位置与q的概率质量函数 $Pq_{i}$ 相同：

即 $\tilde{y}$ 的可微熵可以作为q的熵的近似

这样将 $\tilde{y}$ 的可微熵代替q的熵带入损失函数中，使得损失函数能够实现梯度下降，同时保证模型的精确性。

四、熵模型

熵模型用于在图像压缩中估计量化后的潜在表示的熵（最小码长），在训练阶段，通过用加性均匀噪声替代量化操作，使损失函数可导，从而支持反向传播优化模型；在实际压缩阶段，使用类似 CABAC 的算术编码框架，根据量化值的概率模型生成二进制流。

根据香农信源编码定理，编码比特率的理论下限是符号的熵值。通过学习量化后系数的概率质量函数，算术编码能够将每个符号的平均比特数逼近其熵值，从而最小化传输或存储所需的比特率。

（一）训练阶段

熵模型对图像量化后的潜在表示 $\tilde{y}$ 进行先验建模，对编码特征的分布进行精确预测，为算术编码提供精确的概率估计。

self.prior = tfc.NoisyDeepFactorized(batch_shape=(num_filters,))

假设 $\tilde{y}$ 的各个分量（如不同通道、位置的元素）在统计上相互独立（边缘独立），对 $\tilde{y}$ 的各个分量 $\tilde{y_{i}}$ 并采用分段线性函数来描述其概率密度 $p_{\tilde{y_{i}}}$ ，为 $p_{\tilde{y_{i}}}$ 设定一个可学习的参数向量 $\psi \left ( i \right )$ ，这些参数控制分段线性函数的形状，决定 $p_{\tilde{y_{i}}}$ 的具体形式。其中的元素是从 $p_{\tilde{y_{i}}}$ 中采样得到的，采样密度为每单位长度采样 10 个点，通过使用普通随机梯度下降最大化负的期望似然函数来优化参数 $\psi \left ( i \right )$ ，最终使先验模型适配数据特征：

（二）实际压缩阶段

算术编码是一种无损压缩技术，通过将符号序列映射为 [0,1) 区间内的实数，并用尽可能少的比特表示该实数。其效率依赖于准确的概率模型，压缩率接近符号的熵值。

量化后图像数据q：将q在空间上光栅扫描并在通道上依次迭代处理形成一串数据流，基于训练得到的概率模型，对优化后的 $p_{\tilde{y_{i}}}$ 进行采样得到 $p_{q_{i}}$

对量化后的q，按二进制决策树生成二进制码流，对于每个 $q_{i}$ ，我们首先测试编码的值是否等于分布的众数（即出现概率最高的值）。如果是这样，则 $q_{i}$ 直接编码为一个二进制符号；若不是，则进一步通过另一个二进制判定来决定这个值是大于还是小于众数值，从而实现更精细的编码。

接着，依次测试每个可能的整数值，形成一个二叉决策链，直到确定 $q_{i}$ 的值，或者达到最小 $q_{i,min}$ 或最大 $q_{i,max}$ 值。对于不在 [ $q_{i,min}$ , $q_{i,max}$ ]的值，使用哥伦布码编码其与边界的差值，跳过算术编码。

五、压缩后输出

熵编码的输出是一个压缩后的比特流，这个比特流包含了量化后离散值的编码信息。除了量化值的编码，比特流中还包含图像的尺寸、图像类型和模型参数信息，以便在解码时能够正确恢复图像。

六、重构图像

熵编码的比特流解码回潜在表示 $\hat{y}$

$\hat{y}$ 经过合成变换转换回图像

将转换后的图像的像素范围从[0,1]变回[0,255]，得到最终重构图像 $\hat{x}$

七、训练模型——率失真联合优化

模型通过率失真权衡损失函数进行端到端优化：L=R+λ⋅D

R为编码潜在表示所需的比特率，D为重构图像与原始图像的均方误差（MSE）。

训练过程：

前向传播：图像经过分析变换、量化、熵编码、合成变换，得到重构图像。

损失计算：计算MSE和熵估计，加权求和。

反向传播：通过梯度下降，更新分析变换、合成变换和熵模型的参数。

def call(self, x, training):#计算loss，bpp，mse
    """Computes rate and distortion losses."""
    entropy_model = tfc.ContinuousBatchedEntropyModel(#熵模型
        self.prior, coding_rank=3, compression=False)#指定了编码时的秩（rank）。在图像压缩中，通常选择 3 来表示 RGB 图像的三个颜色通道。
    x = tf.cast(x, self.compute_dtype)  # TODO(jonarchist): Why is this necessary?
    y = self.analysis_transform(x)
    y_hat, bits = entropy_model(y, training=training)#y_hat：编码后的潜在表示；bits：编码过程中使用的比特数，即每个像素的比特数，表示压缩的比特率。
    x_hat = self.synthesis_transform(y_hat)#重建图像 x_hat
    # Total number of bits divided by total number of pixels.
    num_pixels = tf.cast(tf.reduce_prod(tf.shape(x)[:-1]), bits.dtype)#这段代码获取输入图像 x 的形状，并排除最后一个维度（即颜色通道维度）。这为计算图像的空间尺寸（宽度和高度）提供了信息。
    bpp = tf.reduce_sum(bits) / num_pixels#计算图像的总像素数
    # Mean squared error across pixels.
    mse = tf.reduce_mean(tf.math.squared_difference(x, x_hat))#计算原始图像 x 和重建图像 x_hat 之间的像素差的平方。
    mse = tf.cast(mse, bpp.dtype)#求解均方误差
    # The rate-distortion Lagrangian.
    loss = bpp + self.lmbda * mse#计算损失函数
    return loss, bpp, mse

训练过程代码：

def train_step(self, x):
    with tf.GradientTape() as tape:#tf.GradientTape 是 TensorFlow 中用于自动求导的上下文管理器。在 with 块内的所有操作都会被记录下来，以便后续计算梯度。
      loss, bpp, mse = self(x, training=True)#将输入数据 x 传入模型进行前向传播。training=True 表示当前处于训练模式
    variables = self.trainable_variables#self.trainable_variables 获取模型中所有可训练的变量
    gradients = tape.gradient(loss, variables)#tape.gradient(loss, variables) 计算损失函数 loss 相对于可训练变量 variables 的梯度。
    self.optimizer.apply_gradients(zip(gradients, variables))#apply_gradients 方法根据计算得到的梯度 gradients 对可训练变量 variables 进行更新，以最小化损失函数。
    self.loss.update_state(loss)
    self.bpp.update_state(bpp)
    self.mse.update_state(mse)
    return {m.name: m.result() for m in [self.loss, self.bpp, self.mse]}