Source-linked AI summary
CompressAI: a PyTorch library and evaluation platform for end-to-end compression research
Jean Bégaint, Fabien Racapé, Simon Feltman, Akshay Pushparaja
TL;DR
Learned compression research needs tooling for custom operations, model development, and fair evaluation, while PyTorch lacks many compression-specific components. CompressAI provides a PyTorch platform with reimplemented models, pretrained weights, and codec-comparison tools. Its current scope centers on still-picture compression, with video support planned, and runtime comparisons remain difficult to assess fairly.
Problem
PyTorch lacks custom operations needed for end-to-end compression architectures, requiring substantial reimplementation work.
Method
CompressAI provides PyTorch compression operations, evaluation tools, reimplemented learned image-compression models, pretrained weights, and accessible training and inference APIs.
Results
CompressAI reproduces results from the literature and enables evaluation against learned and conventional codecs; it has also been adopted by industrial and academic research communities.
Takeaways & Limitations
CompressAI provides an open research pipeline for training and evaluating neural compression codecs and comparing them with traditional codecs.
Takeaways & Limitations
Current tools and documentation focus mostly on learned image compression, with video-compression support planned for future releases.
Abstract
from arXiv · showhide
This paper presents CompressAI, a platform that provides custom operations, layers, models and tools to research, develop and evaluate end-to-end image and video compression codecs. In particular, CompressAI includes pre-trained models and evaluation tools to compare learned methods with traditional codecs. Multiple models from the state-of-the-art on learned end-to-end compression have thus been reimplemented in PyTorch and trained from scratch. We also report objective comparison results using PSNR and MS-SSIM metrics vs. bit-rate, using the Kodak image dataset as test set. Although this framework currently implements models for still-picture compression, it is intended to be soon extended to the video compression domain.
1 Introduction
ANN-based codecs learn nonlinear transforms for image compression and have also advanced video compression, but the field lacks tooling for further research. CompressAI is presented to address this tooling gap.
- Conventional codecs use block partitioning, linear transforms, prediction, quantization, and entropy coding, whereas ANN-based codecs rely on learned components.
- ANN-based codecs use learned nonlinear transforms and optimize bitstream length alongside reconstructed-image distortion.Distortion can be measured with MSE or MS-SSIM.
- Neural video compression is more challenging because reducing temporal redundancy requires motion estimation, larger networks, and multi-stage training.
- Despite promising ANN-based image and video compression results, the field needs more research and experiments and lacks tooling to facilitate contributions.CompressAI aims to help improve this situation.
2 Motivation
PyTorch lacks the custom compression operations needed to build end-to-end learned codecs from scratch, creating substantial reimplementation work. CompressAI supplies these operations, evaluation tools, reimplemented models, and pretrained weights for a complete research pipeline.
- PyTorch does not ship with custom compression operations such as entropy bottlenecks or entropy coding tools, and these components are mostly absent from its ecosystem.
- CompressAI implements common PyTorch compression operations and evaluation tools for comparing learned methods with traditional codecs.
- CompressAI reimplements state-of-the-art learned image-compression models and provides pretrained weights across multiple bit-rate points and quality metrics.The weights were learned using the Vimeo-90K training dataset and achieve similar performances to reported original-paper numbers.
- CompressAI enables a complete pipeline from training through performance evaluation against learned and conventional codecs.
3 Design
CompressAI follows PyTorch conventions and learned-compression terminology to make codec development and use more accessible. Its high-level APIs support training and inference without requiring extensive prior knowledge.
- CompressAI follows PyTorch design principles and the code structures and conventions of widely used PyTorch libraries.
- The library follows naming conventions from learned data-compression literature to ease the transition from papers to code.
- High-level training and inference APIs require little prior knowledge of learned compression or deep learning.
- Specific compression-domain implementations may still require familiarity with learned-compression terminology.
4 Features
CompressAI provides PyTorch components, pre-trained models, and evaluation tools for developing and comparing end-to-end compression codecs. Its current scope emphasizes learned image compression, with video support planned for future extension.
- Neural-network building blocks: CompressAI implements entropy models, quantization operations, color transforms, and other domain-specific modules for building end-to-end compression networks in PyTorch.
- Neural-network building blocks: Multiple state-of-the-art learned image-compression architectures were re-implemented in PyTorch, with training, fine-tuning, inference, and evaluation supported.
- Scope: The tools and documentation currently focus mainly on learned image compression, while video-compression support is planned.
- Model zoo: Pre-trained models are available through a model zoo for multiple quality settings, bit-rates, and MSE or MS-SSIM metrics.
- Model zoo: Training uses randomly cropped 256 × 256 Vimeo-90K image patches, typically for 4–5M steps, with training lasting approximately 4–10 days depending on the model and hardware.
- Evaluation and benchmarking: CompressAI supplies command-line and Python tools to evaluate learned and conventional codecs, find quality parameters, and compare methods using PSNR and MS-SSIM.
- Evaluation and benchmarking: Runtime comparisons between traditional and ANN-based methods are not currently reported as accurate and fair because their typical hardware execution differs.
5 Evaluation
CompressAI reimplements learned compression models, provides pretrained weights, and evaluates them against original publications and traditional codecs on Kodak. The reproduced results are similar to published results, while learned methods compare favorably with HEVC, AV1, and VVC.
- Model reproduction: CompressAI reimplements factorized-prior, hyperprior, autoregressive, joint-hyperprior, and self-attention compression models.
- Model reproduction: Pretrained weights optimized for MSE are available through the CompressAI API, with MS-SSIM weights being added.
- Comparison with originally published results: Similar average Kodak performance was reproduced relative to the original publications using actual bit-rates from complete encoding and decoding pipelines.Floating-point operations prevent reproducibility across different systems or platforms.
- Learned codecs versus traditional codecs: Learned image-compression methods compare favorably with HEVC and AV1, with the strongest methods competitive with VVC in PSNR at low bit-rates.Learned frameworks can also be optimized for differentiable objective metrics.
- Comparison with originally published results: Figure 1 compares CompressAI and original results using PSNR (RGB) versus bit-rate on Kodak, despite different training data and implementations.
6 Conclusion and future work
CompressAI currently supports still-picture coding with pretrained models and tools for comparison, reproduction, and custom neural codec development. Planned extensions target broader models, perceptual metrics, video compression, interoperability, and continued open development.
- Current platform: CompressAI implements still-picture networks, pretrained weights, and tools for comparing state-of-the-art models with traditional image codecs.
- Current platform: The platform reproduces literature results and lets researchers, developers, and enthusiasts train and evaluate their own neural-network-based codecs.
- Future work: Planned releases add models, perceptual-metric weights, video-compression support, video evaluation modes, compressible motion modules, and improved TorchScript and ONNX compatibility.
- Open development: CompressAI is available under the Apache 2.0 license and will continue to be supported and extended openly on GitHub.
A.1 PSNR on Kodak
PSNR rate-distortion curves on Kodak show learned methods outperforming older image codecs and challenging newer standards.
- JPEG, JPEG 2000, and WebP are largely outperformed by all learned methods on Kodak PSNR rate-distortion curves.
- Hyperprior-based methods can reach similar or better PSNR performance than HEVC/BPG, AV1, and VVC.
A.2 MS-SSIM on Kodak
On Kodak, learned methods significantly outperform traditional codecs under MS-SSIM evaluation, including models fine-tuned for that metric.
- Hyperprior and factorized models from [1] are fine-tuned using the MS-SSIM metric.
- Learned methods significantly outperform traditional methods under MS-SSIM, while MSE-trained models remain competitive or better.
A.3 PSNR on CLIC Mobile (2020)
Figure 5 presents rate-distortion curves for PSNR measured on the CLIC Mobile dataset.
- Figure 5 reports PSNR rate-distortion curves on the CLIC Mobile dataset.
A.4 MS-SSIM on CLIC Mobile (2020)
The paper includes MS-SSIM rate-distortion curves for CLIC Mobile and CLIC Pro, alongside visual comparisons and implementation examples.
- Figure 6 presents MS-SSIM rate-distortion curves measured on the CLIC Mobile dataset.
- Figure 8 presents MS-SSIM rate-distortion curves measured on the CLIC PRO dataset.