CompressAI: a PyTorch library and evaluation platform forend-to-end compression research，我的理解是一个基于图像（视频）压缩的API库。他是建立在 PyTorch 之上的，用于基于深度学习的数据压缩的自定义操作、层和模型，其中包括了基于tensorflow.compression压缩包的部分移植，还包括包括一些用于压缩任务的预训练模型。其中可以直接用到的模型来自以下几篇文献：
- bmshj2018-factorized
Ballé J, Laparra V, Simoncelli E P. End-to-end optimized image compression[C]//ICLR 2017 - bmshj2018-hyperprior
Ballé J, Minnen D, Singh S, et al. Variational image compression with a scale hyperprior[C]//ICLR2018 - mbt2018-mean
Ballé J, Toderici G. Joint autoregressive and hierarchical priors for learned image compression[C]//NIPS 2018. - mbt2018
Minnen D, Ballé J, Toderici G. Joint autoregressive and hierarchical priors for learned image compression[C]//NIPS 2018. - cheng2020-anchor
Cheng Z, Sun H, Takeuchi M, et al. Learned Image Compression With Discretized Gaussian Mixture Likelihoods and Attention Modules[C]//CVPR 2020 - 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简介：简介说明以及使用说明
整个项目基于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 models1、总览文件结构
这里主要看compressai和example两个文件夹
封装了相关的api，其中models可以看到是封装的模型
example下包含了一些代码使用的案例，还有train.py
我这里着重看了train.py，包含了我们设置的一些参数以及训练模型的主要流程，调用这个py文件通过命令的方式进行，在命令中设置数据集地址、训练参数等信息，调用此文件进行训练
首先我们要在特定文件夹下放好数据集，根据官方给的地址是/path/to/my/image/dataset/，我们这里采用我们存放数据集的地址，该数据集文件夹下分为了两个部分，一个train，一个test，这部分可以和train.py的代码中看到，根据字符串截取获取了对应数据集的文件
官方给的命令是
python3 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save
相关命令说明在代码里有
我们将根据我们需要的命令、数据集地址、调用的模型等参数进行修改
- 加载训练命令，如上述图片的数据集、 e p o c h epoch epoch、学习率、 p a t c h s i z e patchsize patchsize、随机种子等
- 对数据集图片进行剪裁
分别读取train训练数据集与test测试数据集的图片，剪裁大小 p a t c h s i z e patchsize patchsize初始值是256，train训练集随机剪裁成 p a t c h s i z e ∗ p a t h s i z e patchsize*pathsize patchsize∗pathsize，test测试数据集中心剪裁成 p a t c h s i z e ∗ p a t h s i z e patchsize*pathsize patchsize∗pathsize - 加载剪裁好的两个数据集，以及训练相关数据如 b a t c h s i z e batchsize batchsize等
- 读取需要训练的模型网络、进行优化器的初始设置
- 训练前先判断是否在之前有保存的节点，有则继续训练，无则重新开始
- 对train训练集以及test测试集分别带入优化器进行梯度下降，开始模型训练
- 损失收敛后训练完成，保存模型以及相应的节点
python -m compressai.utils.update_model --architecture ARCH checkpoint_best_loss.pth.tar
ARCH可换成对应训练的模型如mbt2018
评估训练模型
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...
具体实例的使用我将再写一篇博客进行记录
暑期编程PK赛 得CSDN机械键盘等精美礼品！
