Skip to main content

A collection of tools for neural compression enthusiasts

Project description


LICENSE Build and Test Documentation Status

What's New


NeuralCompression is a Python repository dedicated to research of neural networks that compress data. The repository includes tools such as JAX-based entropy coders, image compression models, video compression models, and metrics for image and video evaluation.

NeuralCompression is alpha software. The project is under active development. The API will change as we make releases, potentially breaking backwards compatibility.


NeuralCompression is a project currently under development. You can install the repository in development mode.

PyPI Installation

First, install PyTorch according to the directions from the PyTorch website. Then, you should be able to run

pip install neuralcompression

to get the latest version from PyPI.

Development Installation

First, clone the repository and navigate to the NeuralCompression root directory and install the package in development mode by running:

pip install --editable ".[dev, docs]"

If you are not interested in matching the test environment, then you only need to apply the second step to install.

Repository Structure

We use a 2-tier repository structure. The neuralcompression package contains a core set of tools for doing neural compression research. Code committed to the core package requires stricter linting, high code quality, and rigorous review. The projects folder contains code for reproducing papers and training baselines. Code in this folder is not linted aggressively, we don't enforce type annotations, and it's okay to omit unit tests.

The 2-tier structure enables rapid iteration and reproduction via code in projects that is built on a backbone of high-quality code in neuralcompression.


  • neuralcompression - base package
    • data - PyTorch data loaders for various data sets
    • entropy_coders - lossless compression algorithms in JAX
      • craystack - an implementation of the rANS algorithm with the craystack API
    • functional - methods for image warping, information cost, flop counting, etc.
    • layers - building blocks for compression models
    • metrics - torchmetrics classes for assessing model performance
    • models - complete compression models


Tutorial Notebooks

This repository also features interactive notebooks detailing different parts of the package, which can be found in the tutorials directory. Existing tutorials are:

  • Walkthrough of the neuralcompression flop counter (view on Colab).
  • Using neuralcompression.metrics and torchmetrics to calculate rate-distortion curves (view on Colab).

Getting Started

For an example of package usage, see the Scale Hyperprior for an example of how to train an image compression model in PyTorch Lightning. See DVC for a video compression example.


Please read our CONTRIBUTING guide and our CODE_OF_CONDUCT prior to submitting a pull request.

We test all pull requests. We rely on this for reviews, so please make sure any new code is tested. Tests for neuralcompression go in the tests folder in the root of the repository. Tests for individual projects go in those projects' own tests folder.

We use black for formatting, isort for import sorting, flake8 for linting, and mypy for type checking.


NeuralCompression is MIT licensed, as found in the LICENSE file.


If you find NeuralCompression useful in your work, feel free to cite

    author={Matthew Muckley and Jordan Juravsky and Daniel Severo and Mannat Singh and Quentin Duval and Karen Ullrich},

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

neuralcompression-0.2.1.tar.gz (1.7 MB view hashes)

Uploaded source

Built Distribution

neuralcompression-0.2.1-py3-none-any.whl (104.0 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page