CompressAI: A PyTorch Toolkit for Learned Image and Video Compression Research

The gap CompressAI fills between learned codecs and classical benchmarks

Learned image compression papers usually ship as a model definition plus a training script. What they rarely ship is a way to put the resulting bitstream next to BPG or VTM on the same images with the same metric code. CompressAI's stated purpose is that second half: it describes itself as a PyTorch library and evaluation platform for end-to-end compression research, and it provides evaluation scripts to compare learned models against classical image and video codecs. The audience is therefore narrow and identifiable. It is for researchers who already work in PyTorch, who need entropy model layers and pre-trained checkpoints to build on, and who need a bench command that produces comparable rate-distortion numbers. It is not aimed at application developers who want to drop a codec into a product. The README's own framing, research, is the honest description of scope.

What is actually inside: layers, a partial TFC port, and a model zoo

The README lists four things the library currently provides. First, custom operations, layers and models for deep learning based data compression. Second, a partial port of the official TensorFlow compression library, which means the entropy model machinery familiar from that ecosystem has a PyTorch counterpart here, though the word partial matters: this is not a complete translation. Third, pre-trained end-to-end compression models for learned image compression, documented on a model zoo page. Fourth, the evaluation scripts already mentioned. The repository also carries a PSNR performance plot on Kodak in its assets, which is the kind of evidence the project uses to position its models. Note what the list does not contain: no claim of production bitstream stability, no format specification. The deliverable is code and checkpoints for experimentation.

Training and evaluation flow through the examples directory

The mechanism visible in the README is a set of scripts under examples/ rather than a single framework entry point. Training uses examples/train.py with a rate-distortion loss, and the README states you can swap in your own model implemented within CompressAI and run the same pipeline. The invocation shown is python3 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save. Video training mirrors this with examples/train_video.py. Multi-GPU runs go through torchrun, and the README is explicit that --batch-size is per process, so two GPUs means the effective global batch size is 2 x --batch-size. That note is worth internalising before you compare learning rates across single and multi-GPU runs. Evaluation is separate and lives in the package namespace: python3 -m compressai.utils.eval_model checkpoint /path/to/images/folder/ -a $ARCH -p $MODEL_CHECKPOINT for your own checkpoints, or the pretrained subcommand with -q $QUALITY_LEVELS for the shipped models. Plotting is a third step via compressai.utils.plot. The data flow is thus: dataset in a custom ImageFolder layout, train script producing a checkpoint, eval_model producing metrics, plot producing figures.

Installation paths and the packaging constraints the README admits

The simplest route is pip install compressai, with the README noting wheels are available for Linux and MacOS. Source installs need a C++17 compiler, pip 19.0 or newer, and the packages listed in setup.py. The documented sequence clones into a directory named compressai, then runs pip install -U pip && pip install -e . inside a virtual environment. Two extras exist: pip install -e '.[dev]' for testing, linting and docs, and pip install -e '.[tutorials]' for the notebook dependencies. The README states CompressAI supports Python 3.8+ and PyTorch 1.7+. It also states plainly that Docker images will be released in the future and that Conda environments are not officially supported. If your lab standardises on Conda, that sentence is the one to read twice. Multi-GPU support is flagged as experimental, which is a meaningful caveat for anyone planning distributed training as a default rather than an occasional run.

Benchmarking against BPG, VTM and x265, and the external dependencies that implies

The classical comparison side runs through compressai.utils.bench. The README shows python3 -m compressai.utils.bench --help, then bpg --help and vtm --help as subcommands. For video, compressai.utils.video.bench is invoked with x265 or VTM, and video evaluation currently covers ssf2020 only, which the README states directly: CompressAI only includes ssf2020 for now. That single-model video scope is the clearest limitation in the material. The bench commands wrap external codecs rather than reimplementing them, so the practical constraint is that BPG, VTM or x265 must be available on your system for those comparisons to run. The README does not document how those binaries are located or configured. Anyone planning a benchmark suite should treat binary provisioning as their own problem and confirm it before scheduling a large evaluation job.

Where CompressAI is the wrong choice, and what TensorFlow Compression does differently

The most direct alternative named in the README's related links is the TensorFlow compression library by Ballé et al., which CompressAI partially ports. The difference in approach is ecosystem, not algorithm: TensorFlow Compression is built around TensorFlow ops and its own training conventions, while CompressAI reimplements the relevant layers in PyTorch and adds an evaluation harness that spans classical codecs. If your existing training stack, data pipeline and checkpoint tooling are TensorFlow, porting to CompressAI buys you the bench scripts at the cost of a framework migration. If they are PyTorch, the reverse is true. Beyond that comparison, CompressAI is the wrong tool in three concrete cases. You need a deployable codec with a frozen bitstream: nothing in the README promises that. You need broad video codec coverage: only ssf2020 is included. You need Conda packaging or a maintained Docker image today: the README defers both. The library is also, by its own description, a research platform, so API churn between versions should be assumed rather than treated as a surprise.

Licence, citation obligations and the cost of staying current

CompressAI is licensed under the BSD 3-Clause Clear License, stated in the README and in the repository's LICENSE file. That is a permissive licence with an explicit patent-related clause, but the details are for you and your organisation's counsel to read in the LICENSE text rather than for a summary here. Two citation obligations are documented. The project itself should be cited via the begaint2020compressai arXiv entry, and the README adds a separate citation for variable bitrate models, kamisli2024dcc_vbrlic from the 2024 Data Compression Conference. If you publish results from the variable-rate models, that second citation applies. On maintenance cost: the repository is not archived and the last push recorded is 2026-07-04, so it is active. No releases were retrieved for this review, so I cannot describe a version cadence. The recurring costs that are visible in the material are the C++17 toolchain for source builds, the external codec binaries for benchmarks, and the per-process batch size arithmetic whenever you change GPU count. None of those is large, but each one silently changes your results if you get it wrong.

Editorial conclusion

Adopt CompressAI if you are doing learned compression research in PyTorch and want pre-trained models plus evaluation scripts that can be compared against BPG, VTM or x265 in the same harness. Do not adopt it if you need a shipping codec, a stable API surface, or Conda-based packaging: the README states Conda environments are not officially supported and Docker images are a future item. Before committing, verify the model zoo page for the architecture you intend to use, check that your Python is 3.8+ and your PyTorch 1.7+, and run the benchmark command for your target classical codec to confirm the external binary is present on your machine.

Official sources

1. InterDigitalInc/CompressAI on GitHub
2. Issues
3. License: BSD-3-Clause-Clear
4. Project website
5. README