FlashGMM is a fast implementation of Gaussian Mixture Model (GMM) for image compression in PyTorch. The codes are based on the CompressAI library.
Our paper is available at arxiv. Please feel free to contact me (octachoron(at)suou.waseda.jp) or post an issue if you have any inquiries regarding to this work.
The Gaussian Mixture Model (GMM) by Cheng et al. 2020 is a powerful entropy model for learned image compression. However, the original implementation is computationally expensive, making it impractical for real-world applications.
Our implementation accelerates the GMM-based entropy model by leveraging parallel computing capabilities of modern GPUs. We achieve this by (1) installing a new search-based algorithm, (2) numerical approximation of the cumulative distribution function (CDF), and (3) leveraging SIMD parallelism in CPU.
Our impelmentation is fully compatible with the CompressAI's implementation. That is, you can apply FlashGMM to any models trained with CompressAI's GMM entropy model without any modifications.
We currently support the following models:
- Cheng2020AnchorCheckerboardGMMv2 (He et al. CVPR 2021) [weights]
- ELIC2022GMM (He et al. CVPR 2022)
We tested FlashGMM on Python 3.10, Pytorch 2.2.2, CUDA 12.1 and GCC 11.4.0. We recommend using Docker container pytorch/pytorch:2.2.2-cuda12.1-cudnn8-devel for easy setup.
First, clone the repository and install the required packages:
git clone https://github.com/tokkiwa/FlashGMM.git
cd FlashGMM
pip install -r requirements.txtTo compile the C++ extensions, run:
python setup.py installThe core implementation of our FlashGMM is in compressai/cpp_exts/rans/rans_interface.cpp and compressai/entropy_models/entropy_models.py. The model code used in the experiment can be found in compressai.models.ckbd_gmm.Cheng2020AnchorCheckerboardGMMv2.
To train the model, you can use the provided training script. For example:
python3 train_ckbd_gmm.py --cuda -d /path/to/dataset
--N 192 --K 4 --lambda 0.0035 --epochs 150 --learning-rate 1e-4 --lr_epoch 120 140 145 --batch-size 16 \
--save_path /path/to/log --patch-size 256 256 \
--kodak_path /path/to/kodakFor evaluation, run
python3 eval.py --data /path/to/kodak --cuda \
--checkpoint /path/to/checkpoint --real\You can set environment variable APPROX_MODE and USE_SIMD to control the behavior of FlashGMM. APPROX_MODE can be set to 0, 1, or 2 to choose the approximation method. 0(default) corresponds to Polya's approximation, 1 to logistic approximation, and 2 to A&S approximation.
USE_SIMD can be set to 0 or 1 to disable or enable SIMD acceleration. It is enabled by default.
The following is the original Readme of CompressAI:
CompressAI (compress-ay) is a PyTorch library and evaluation platform for end-to-end compression research.
CompressAI currently provides:
- custom operations, layers and models for deep learning based data compression
- a partial port of the official TensorFlow compression library
- pre-trained end-to-end compression models for learned image compression
- evaluation scripts to compare learned models against classical image/video compression codecs
Note: Multi-GPU support is now experimental.
CompressAI supports python 3.8+ and PyTorch 1.7+.
pip:
pip install compressaiNote: wheels are available for Linux and MacOS.
From source:
A C++17 compiler, a recent version of pip (19.0+), and common python packages
are also required (see setup.py for the full list).
To get started locally and install the development version of CompressAI, run the following commands in a virtual environment:
git clone https://github.com/InterDigitalInc/CompressAI compressai
cd compressai
pip install -U pip && pip install -e .For a custom installation, you can also run one of the following commands:
pip install -e '.[dev]': install the packages required for development (testing, linting, docs)pip install -e '.[tutorials]': install the packages required for the tutorials (notebooks)pip install -e '.[all]': install all the optional packages
Note: Docker images will be released in the future. Conda environments are not officially supported.
Script and notebook examples can be found in the examples/ directory.
To encode/decode images with the provided pre-trained models, run the
codec.py example:
python3 examples/codec.py --helpAn examplary training script with a rate-distortion loss is provided in
examples/train.py. You can replace the model used in the training script
with your own model implemented within CompressAI, and then run the script for a
simple training pipeline:
python3 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --saveNote: the training example uses a custom ImageFolder structure.
A jupyter notebook illustrating the usage of a pre-trained model for learned image
compression is also provided in the examples directory:
pip install -U ipython jupyter ipywidgets matplotlib
jupyter notebook examples/To evaluate a trained model on your own dataset, CompressAI provides an evaluation script:
python3 -m compressai.utils.eval_model checkpoint /path/to/images/folder/ -a $ARCH -p $MODEL_CHECKPOINT...To evaluate provided pre-trained models:
python3 -m compressai.utils.eval_model pretrained /path/to/images/folder/ -a $ARCH -q $QUALITY_LEVELS...To plot results from bench/eval_model simulations (requires matplotlib by default):
python3 -m compressai.utils.plot --helpTo evaluate traditional codecs:
python3 -m compressai.utils.bench --help
python3 -m compressai.utils.bench bpg --help
python3 -m compressai.utils.bench vtm --helpFor video, similar tests can be run, CompressAI only includes ssf2020 for now:
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 --helpRun tests with pytest:
pytest -sx --cov=compressai --cov-append --cov-report term-missing testsSlow tests can be skipped with the -m "not slow" option.
CompressAI is licensed under the BSD 3-Clause Clear License
We welcome feedback and contributions. Please open a GitHub issue to report bugs, request enhancements or if you have any questions.
Before contributing, please read the CONTRIBUTING.md file.
- Jean Bégaint, Fabien Racapé, Simon Feltman and Hyomin Choi, InterDigital AI Lab.
If you use this project, please cite the relevant original publications for the models and datasets, and cite this project as:
@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},
}
For any work related to the variable bitrate models, please cite
@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},
}
- Tensorflow compression library by Ballé et al.: https://github.com/tensorflow/compression
- Range Asymmetric Numeral System code from Fabian 'ryg' Giesen: https://github.com/rygorous/ryg_rans
- BPG image format by Fabrice Bellard: https://bellard.org/bpg
- HEVC HM reference software: https://hevc.hhi.fraunhofer.de
- VVC VTM reference software: https://vcgit.hhi.fraunhofer.de/jvet/VVCSoftware_VTM
- AOM AV1 reference software: https://aomedia.googlesource.com/aom
- Z. Cheng et al. 2020: https://github.com/ZhengxueCheng/Learned-Image-Compression-with-GMM-and-Attention
- Kodak image dataset: http://r0k.us/graphics/kodak/