CompressAI (compress-ay) 是一个用于端到端压缩研究的 PyTorch 库和评估平台。

目前，CompressAI 提供：

• 用于基于深度学习的数据压缩的自定义操作、层和模型
• 官方 TensorFlow compression 库的部分移植
• 预训练的端到端压缩模型，用于学习型图像压缩
• 用于将学习型模型与经典图像/视频压缩编解码器进行比较的评估脚本

注意：多 GPU 支持目前仍处于实验阶段。

安装

CompressAI 支持 Python 3.8 及以上版本和 PyTorch 1.7 及以上版本。

通过 pip：

pip install compressai

注意：适用于 Linux 和 macOS 的预编译二进制包已提供。

从源代码安装：

还需要 C++17 编译器、较新版本的 pip（19.0 及以上）以及常见的 Python 包（完整列表请参见 setup.py 文件）。

要在本地开始使用并安装 CompressAI 的开发版本，请在 虚拟环境 中运行以下命令：

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

对于自定义安装，您还可以运行以下命令之一：

• pip install -e '.[dev]'：安装开发所需的包（测试、代码风格检查、文档等）
• pip install -e '.[tutorials]'：安装教程所需的包（Notebook 等）

注意：未来将发布 Docker 镜像。Conda 环境目前未被官方支持。

文档

• 安装指南
• CompressAI API
• 训练您自己的模型
• 可用模型列表（模型库）

使用

示例

脚本和 Notebook 示例可在 examples/ 目录中找到。

要使用提供的预训练模型对图像进行编码/解码，请运行 codec.py 示例：

python3 examples/codec.py --help

examples/train.py 中提供了一个使用率失真损失的示例训练脚本。您可以将该脚本中使用的模型替换为您在 CompressAI 中实现的自定义模型，然后运行该脚本来完成一个简单的训练流程：

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

要在 2 个 GPU 上使用分布式数据并行运行相同的训练脚本，请使用 torchrun：

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save

注意：--batch-size 是每个进程的批次大小，因此在使用 2 个 GPU 时，实际的全局批次大小为 2 x --batch-size。

注意：该训练示例使用了自定义的 ImageFolder 数据结构。

对于视频训练，可以以相同的方式启动 examples/train_video.py：

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 examples/train_video.py -d /path/to/my/video/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save

注意：对于 examples/train_video.py，--batch-size 同样是每个进程的批次大小，因此在使用 2 个 GPU 时，实际的全局批次大小为 2 x --batch-size。

examples 目录中还提供了一个 Jupyter Notebook，演示如何使用预训练模型进行学习型图像压缩：

pip install -U ipython jupyter ipywidgets matplotlib
jupyter notebook examples/

评估

要对您自己的数据集上的训练模型进行评估，CompressAI 提供了一个评估脚本：

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

要评估提供的预训练模型：

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

要绘制 bench/eval_model 模拟的结果（默认需要 matplotlib）：

python3 -m compressai.utils.plot --help

要评估传统编解码器：

python3 -m compressai.utils.bench --help
python3 -m compressai.utils.bench bpg --help
python3 -m compressai.utils.bench vtm --help

对于视频，也可以进行类似的测试，目前 CompressAI 仅包含 ssf2020：

python3 -m compressai.utils.video.eval_model checkpoint /path/to/video/folder/ -a ssf2020 -p $MODEL_CHECKPOINT...
python3 -m compressai.utils.video.eval_model pretrained /path/to/video/folder/ -a ssf2020 -q $QUALITY_LEVELS...
python3 -m compressai.utils.video.bench x265 --help
python3 -m compressai.utils.video.bench VTM --help
python3 -m compressai.utils.video.plot --help

测试

使用 pytest 运行测试：

pytest -sx --cov=compressai --cov-append --cov-report term-missing tests

可以通过 -m "not slow" 选项跳过耗时的测试。

许可证

CompressAI 采用 BSD 3-Clause Clear License 许可。

贡献

我们欢迎反馈和贡献。如需报告 bug、请求功能增强或有任何疑问，请在 GitHub 上提交 issue。

在贡献之前，请阅读 CONTRIBUTING.md 文件。

作者

• Jean Bégaint、Fabien Racapé、Simon Feltman 和 Hyomin Choi，InterDigital AI 实验室。

引用

如果您使用本项目，请引用相关模型和数据集的原始论文，并按以下格式引用本项目：

@article{begaint2020compressai,
	title={CompressAI: a PyTorch library and evaluation platform for end-to-end compression research},
	author={B{\'e}gaint, Jean and Racap{\'e}, Fabien and Feltman, Simon and Pushparaja, Akshay},
	year={2020},
	journal={arXiv preprint arXiv:2011.03029},
}

对于与可变比特率模型相关的任何工作，请引用：

@article{kamisli2024dcc_vbrlic,
	title={Variable-Rate Learned Image Compression with Multi-Objective Optimization and Quantization-Reconstruction Offsets},
	author={Kamisli, Fatih and Racap{\'e}, Fabien and Choi, Hyomin},
	year={2024},
	booktitle={2024 Data Compression Conference (DCC)},
	eprint={2402.18930},
}

相关链接

• Ballé 等人 的 TensorFlow 压缩库：https://github.com/tensorflow/compression
• Fabian 'ryg' Giesen 的范围非对称数制编码：https://github.com/rygorous/ryg_rans
• Fabrice Bellard 的 BPG 图像格式：https://bellard.org/bpg
• HEVC HM 参考软件：https://hevc.hhi.fraunhofer.de
• VVC VTM 参考软件：https://vcgit.hhi.fraunhofer.de/jvet/VVCSoftware_VTM
• AOM AV1 参考软件：https://aomedia.googlesource.com/aom
• Z. Cheng 等人，2020 年：https://github.com/ZhengxueCheng/Learned-Image-Compression-with-GMM-and-Attention
• 柯达图像数据集：http://r0k.us/graphics/kodak/

CompressAI 快速上手指南

CompressAI 是一个基于 PyTorch 的端到端压缩研究与评估平台，提供深度学习数据压缩所需的自定义操作、层、模型以及预训练模型。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux 或 MacOS（官方提供 wheel 包支持）。
• Python 版本：3.8 及以上。
• PyTorch 版本：1.7 及以上。
• 编译器：若从源码安装，需要支持 C++17 的编译器。
• 其他依赖：较新版本的 pip (19.0+)。

注意：目前多 GPU 支持处于实验阶段；Conda 环境非官方支持。

安装步骤

方式一：使用 pip 安装（推荐）

这是最简单的安装方式，适用于大多数用户。

pip install compressai

提示：国内用户如遇下载缓慢，可添加清华或阿里镜像源加速： pip install compressai -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装（开发版）

如果您需要修改源码或参与开发，建议使用虚拟环境进行安装：

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

如需安装开发工具链（测试、文档等）或教程依赖（Notebook），可分别运行：

• pip install -e '.[dev]'
• pip install -e '.[tutorials]'

基本使用

1. 图像编解码示例

使用预训练模型对图像进行编码和解码，可以直接运行自带的 codec.py 脚本查看帮助信息：

python3 examples/codec.py --help

2. 模型训练示例

CompressAI 提供了基于率失真损失（rate-distortion loss）的训练脚本。以下命令演示了如何在一个自定义图像数据集上启动训练：

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

参数说明：

• -d: 图像数据集路径（需符合 ImageFolder 结构）。
• --epochs: 训练轮数。
• -lr: 学习率。
• --batch-size: 单进程批次大小。
• --cuda: 启用 GPU 加速。
• --save: 保存检查点。

多 GPU 训练： 若需使用多张显卡（例如 2 张），请使用 torchrun：

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save

注意：多卡模式下，全局有效 batch size = GPU 数量 × --batch-size。

3. 模型评估示例

评估已训练的模型或预训练模型在指定数据集上的性能：

评估自定义模型检查点：

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

评估官方预训练模型：

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

绘制评估结果图表：

python3 -m compressai.utils.plot --help

更多详细用法（如视频压缩、传统编解码器对比）请参考官方文档或运行对应脚本的 --help 命令。