CompressAI — PyTorch Library for Data Compression Models

CompressAI is an open-source library and platform for evaluating data compression models

Image for post
Figure 2: Traditional and learned image codecs compared on the Kodek dataset. Source: Github

The tool is based on the PyTorch deep learning framework. The task of data compression is to minimize the size of the data and at the same time preserve information as much as possible.

More about the library

Currently, CompressAI provides:

  • Custom operations, layers, and neural network models for data compression;
  • Part of the functionality from the official TensorFlow library for compression;
  • Pretrained end-to-end compression models for image compression tasks;
  • Scripts for evaluating and comparing trained models with classical methods of compressing photo and video data

The current version of the library does not support multi-GPU training. However, the developers plan to add this functionality in future versions of the library.

Image for post
Lambda values used for training the networks are different bit-rates (quality setting from 1 to 8), for the MSE and MS-SSIM metrics. Source: Arxiv

Available DL Models

CompressAI’s neural networks were trained in 4–5 million iterations on 256-by-256 image patches that were accidentally extracted and cropped from images in the Vimeo90K dataset. The models were trained with a batch size of 16 or 32. On average, to achieve state-of-the-art quality, the models were trained from one to two weeks, depending on the GPU used.

Multiple architectures from the state-of-the-art on learned image compression have been re-implemented in PyTorch with the domain specific modules and layers provided by CompressAI. There are 6 architectures currently available in CompressAI:

Image for post
Training, fine-tuning, inference, and evaluation of the models listed in this table are fully supported. Source: Arxiv

One of the most important features of CompressAI is the ability to easily implement deep neural networks for end-to-end compression. Several domain-specific layers, operations, and modules have been implemented on top of PyTorch, such as entropy models, quantization operations, color transforms.

Future work

Currently, CompressAI tools and documentation mostly focus on learned image compression and will soon add support for video compression. However, other end-to-end compression pipelines could be built using CompressAI, like the compression of 3D maps or deep features for example.

The developers have also published tutorials on working with the library in the official repository on GitHub.

Original Paper

Interested in Deep Learning?

If you found this article helpful, click the💚 or 👏 button below or share the article on Facebook so your friends can benefit from it too.

Bioinformatician at Oncobox Inc. (@oncobox). Research Associate at Moscow Institute of Physics and Technology (@mipt_eng).

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store