[Paper Review] CompressAI: a PyTorch library and evaluation platform for end-to-end compression research

Jean Bégaint, Fabien Racapé+2 more|arXiv (Cornell University)|Nov 5, 2020

Advanced Data Compression Techniques22 references235 citations

TL;DR

CompressAI provides a PyTorch-based library, pre-trained models, and evaluation tools to research, train, and benchmark end-to-end image compression codecs against traditional codecs.

ABSTRACT

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.

Motivation & Objective

• Provide a PyTorch-based platform to research and develop end-to-end image (and future video) compression codecs.
• Re-implement state-of-the-art learned image compression models with domain-specific components.
• Offer trained weights and a model zoo to reproduce and extend published results.
• Facilitate objective benchmarking of learned codecs against traditional codecs using PSNR and MS-SSIM.
• Enable end-to-end training and evaluation pipelines from training to performance reporting.

Proposed method

• Implement domain-specific PyTorch layers and operations for compression (entropy models, quantization, color transforms).
• Re-implement state-of-the-art learned image compression models in PyTorch and provide pre-trained weights.
• Provide a model zoo with downloadable pre-trained networks for various quality/bit-rate settings.
• Offer training and evaluation pipelines including loss functions that combine distortion and rate terms.
• Provide utilities and a CLI for evaluating models on datasets and benchmarking against traditional codecs.
• Benchmark performance on Kodak/CLIC datasets and report PSNR and MS-SSIM across bit-rates.

Experimental results

Research questions

• RQ1 How close can re-implemented learned image compression models in CompressAI reproduce published rate-distortion results?
• RQ2 How do learned end-to-end codecs compare to traditional codecs (JPEG, JPEG2000, WebP, HEVC, AV1, VVC) in PSNR and MS-SSIM at various bit-rates?
• RQ3 Can a PyTorch-based toolkit streamline training, evaluation, and comparison of learned codecs against conventional standards?
• RQ4 What is the feasibility and trajectory for extending end-to-end compression tooling from images to video?
• RQ5 What training configurations (losses, bit-depths, datasets) yield state-of-the-art performance within CompressAI?

Key findings

• Re-implemented models (factorized prior, hyperprior, Gaussian mixture, autoregressive/hyperprior combinations) reproduce results near those reported in original works when trained from scratch.
• Pre-trained weights enable close replication of reported results and rapid fine-tuning for different metrics or bit-rates.
• Learned codecs significantly outperform traditional methods (e.g., JPEG, JPEG2000, WebP) on PSNR and MS-SSIM benchmarks on Kodak; competitive with modern standards like HEVC, AV1, and approaching VVC at low bit-rates.
• The platform demonstrates that end-to-end learned codecs can be trained and evaluated within a unified PyTorch-based workflow, with reproducibility demonstrated on standard datasets.
• Public adoption of CompressAI is evidenced by its use in industry and academic research, and its influence on DNN-based video coding discussions within standardization groups.