CompressAI：基于pytorch的图像压缩框架使用

前言

CompressAI: a PyTorch library and evaluation platform forend-to-end compression research，我的理解是一个基于图像（视频）压缩的API库。他是建立在 PyTorch 之上的，用于基于深度学习的数据压缩的自定义操作、层和模型，其中包括了基于tensorflow.compression压缩包的部分移植，还包括包括一些用于压缩任务的预训练模型。其中可以直接用到的模型来自以下几篇文献：

1. bmshj2018-factorized
   Ballé J, Laparra V, Simoncelli E P. End-to-end optimized image compression[C]//ICLR 2017
2. bmshj2018-hyperprior
   Ballé J, Minnen D, Singh S, et al. Variational image compression with a scale hyperprior[C]//ICLR2018
3. mbt2018-mean
   Ballé J, Toderici G. Joint autoregressive and hierarchical priors for learned image compression[C]//NIPS 2018.
4. mbt2018
   Minnen D, Ballé J, Toderici G. Joint autoregressive and hierarchical priors for learned image compression[C]//NIPS 2018.
5. cheng2020-anchor
   Cheng Z, Sun H, Takeuchi M, et al. Learned Image Compression With Discretized Gaussian Mixture Likelihoods and Attention Modules[C]//CVPR 2020
6. cheng2020-attn
   Cheng Z, Sun H, Takeuchi M, et al. Learned Image Compression With Discretized Gaussian Mixture Likelihoods and Attention Modules[C]//CVPR 2020

注意： bmshj2018-factorized代码里使用的熵编码方法是Variational image compression with a scale hyperprior提出的全分解方法。官方的tensorflow库里也改了的。

同时提供了与原作者实验对比的测试数据的性能值

并且可以与传统的算法比较效果

相关地址

github：github地址
compressAi简介：简介说明以及使用说明

环境安装

1、pytorch环境

整个项目基于pytorch环境，安装过程可以参考pytorch-gpu版本安装

2、compressai安装

（1）普通版本
pip install compressai

（2）开发版本

git clone https://github.com/InterDigitalInc/CompressAI compressai
cd compressai
pip install -U pip && pip install -e .

安装成功进行验证

python
import compressai

这样一来可以直接在对应环境中引入compressai进行调用

import compressai.models as models

1、总览文件结构

这里主要看compressai和example两个文件夹

1、compressai

封装了相关的api，其中models可以看到是封装的模型

2、example

example下包含了一些代码使用的案例，还有train.py
我这里着重看了train.py，包含了我们设置的一些参数以及训练模型的主要流程，调用这个py文件通过命令的方式进行，在命令中设置数据集地址、训练参数等信息，调用此文件进行训练

2、使用方法

1、准备数据集

首先我们要在特定文件夹下放好数据集，根据官方给的地址是/path/to/my/image/dataset/，我们这里采用我们存放数据集的地址，该数据集文件夹下分为了两个部分，一个train，一个test，这部分可以和train.py的代码中看到，根据字符串截取获取了对应数据集的文件

2、通过命令行进行调用train.py

官方给的命令是

python3 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save

相关命令说明在代码里有


我们将根据我们需要的命令、数据集地址、调用的模型等参数进行修改

整理训练过程

1. 加载训练命令，如上述图片的数据集、$epoch$、学习率、$patchsize$、随机种子等
2. 对数据集图片进行剪裁
   分别读取train训练数据集与test测试数据集的图片，剪裁大小$patchsize$初始值是256，train训练集随机剪裁成$patchsize\ast pathsize$，test测试数据集中心剪裁成$patchsize\ast pathsize$
3. 加载剪裁好的两个数据集，以及训练相关数据如$batchsize$等
4. 读取需要训练的模型网络、进行优化器的初始设置
5. 训练前先判断是否在之前有保存的节点，有则继续训练，无则重新开始
6. 对train训练集以及test测试集分别带入优化器进行梯度下降，开始模型训练
7. 损失收敛后训练完成，保存模型以及相应的节点

3、更新CDF保证熵编码的正常运行（模型地址与1保持一致）

python -m compressai.utils.update_model --architecture ARCH checkpoint_best_loss.pth.tar

ARCH可换成对应训练的模型如mbt2018

4、评价模型，获取性能值

评估训练模型

python3 -m compressai.utils.eval_model checkpoint /path/to/images/folder/ -a $ARCH -p $MODEL_CHECKPOINT...

评估预训练模型

ython3 -m compressai.utils.eval_model pretrained /path/to/images/folder/ -a $ARCH -q $QUALITY_LEVELS...

具体实例的使用我将再写一篇博客进行记录