VARIATIONAL IMAGE COMPRESSION WITH A SCALE HYPERPRIOR文献实验复现

前言

这篇文章是在END-TO-END OPTIMIZED IMAGE COMPRESSION文献基础上进行的改进，主要是加入了超先验网络对边信息进行了处理。相关环境配置与基础可以参考END-TO-END OPTIMIZED IMAGE COMPRESSION
github地址：github

1、相关命令

（1）训练

python bmshj2018.py -V train

同样，我这里方便调试，加入了launch.json

{
    
  // 使用 IntelliSense 了解相关属性。 
  // 悬停以查看现有属性的描述。
  // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
  "version": "0.2.0",
  "configurations": [
    {
    
      "name": "Python: 当前文件",
      "type": "python",
      "request": "launch",
      "program": "${file}",
      "console": "integratedTerminal",
      "justMyCode": true,

      // bmshj2018
      "args": ["--verbose","train"] // 训练


    
    } 
  ]
}

2、数据集

  if args.train_glob:
    train_dataset = get_custom_dataset("train", args)
    validation_dataset = get_custom_dataset("validation", args)
  else:
    train_dataset = get_dataset("clic", "train", args)
    validation_dataset = get_dataset("clic", "validation", args)
  validation_dataset = validation_dataset.take(args.max_validation_steps)

给定数据集地址的情况下使用对应数据集，否则使用默认CLIC数据集

3、整理训练流程

整个代码其实主要还是在bls2017.py上进行的补充，所以整个流转流程是和bls2017的过程一样，这里我简单整理一下这篇文献的训练过程

1. 实例化对象
   init中初始化一些参数（LocationScaleIndexedEntropyModel、ContinuousBatchedEntropyModel两个熵模型中所需参数，对应的模型说明可以见tfc熵模型，具体的参数有scale_fn以及给定的num_scales）
   初始化非线性分析变换、非线性综合变换、超先验非线性分析变换、超先验非线性综合变换、添加均匀噪声、获得先验概率等
   call中通过两个熵模型计算
   （1）x->y（非线性分析变换）
   （2）y->z（超先验非线性分析变换）
   （3）z->z_hat、边信息bit数（ContinuousBatchedEntropyModel熵模型）
   （4）y->y_hat、比特数（LocationScaleIndexedEntropyModel熵模型）
   （5）计算bpp（这里bpp=bit/px，bit数等于上面两项bit数相加）、mse、loss
   其中非线性分析变换即cnn卷积的过程，其中有一步补零的过程是为了输入与输出图片尺寸相等

2. model.compile
   通过model.compile配置训练方法，使用优化器进行梯度下降、算bpp、mse、lose的加权平均

3. 过滤剪裁数据集
   通过参数查看是否给定了数据集路径
   给定数据集路径：直接剪裁成256x256（统一剪裁成256*256送入网络训练，后面压缩的图片不会改变大小）
   未给定数据集路径：用CLIC数据集，过滤出图片大小大于256x256的三通的图片，然后进行剪裁
   分出训练数据集与验证数据集

4. model.fit
   传入训练数据集，设置相关参数（epoch等）进行训练
   通过retval = super().fit(*args, **kwargs)进入model的train_step(不得不说这里封装的太严实了…但从单文件的代码根本找不到怎么跳进去的)

5. train_step
   self.trainable_variables获取变量集，通过传入变量集与定义的损失函数，进行前向传播与反向误差传播更新参数。然后更新loss, bpp, mse

  def train_step(self, x):
    with tf.GradientTape() as tape:
      loss, bpp, mse = self(x, training=True)
    variables = self.trainable_variables
    gradients = tape.gradient(loss, variables)
    self.optimizer.apply_gradients(zip(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]}

6. 接下来就是重复训练的过程，直到到终止条件（比如epoch达到10000）

4、压缩一张图片

python bmshj2018.py [options] compress original.png compressed.tfci

这里通过launch.json调试

"args": ["--verbose","compress","./models/kodak/kodim01.png", "./models/kodim01.tfci"] // 压缩一张图

引入你自己的path