torchac 0.8.9

Fast Arithmetic Coding for PyTorch

torchac: Fast Arithmetic Coding for PyTorch

About

This is a simplified version of the arithmetic coder we used in the neural compression paper "Practical Full Resolution Learned Lossless Image Compression", which lives in the L3C-Pytorch repo. In particular, we removed the L3C-specific parts, which relied on CUDA compliations and were tricky to get going.

The implementation is based on this blog post, meaning that we implement arithmetic coding. While it could be further optimized, it is already much faster than doing the equivalent thing in pure-Python (because of all the bit-shifts etc.). In L3C, Encoding an entire 512 x 512 image happens in 0.202s (see Appendix A in the paper).

What torchac is

• A simple library to encode a stream of symbols into a bitstream given the cumulative distribution of the symbols.
• The number of possible symbols must be finite.

What torchac is not

• We do not provide classes to learn or represent probability/cumulative distributions. These have to be provided by you.

HowTo

Set up conda environment

This library has been tested with

• PyTorch 1.7
• Python 3.8

And that's all you need. Other versions also may work. If you don't have an environment setup, you can make one with conda:

# We use Python 3.8, other version may be supported.
conda create --name <YOUR_ENV_NAME> python==3.8

conda activate <YOUR_ENV_NAME>

# Installing pytorch
Find conda command for your system: https://pytorch.org

Test installation

To (optionally) test your installation, you need pytest:

# If you don't have pytest
pip install pytest

# Run tests
python -m pytest test.py -s

Output should end in something like:

===== 5 passed, 2 warnings in 0.95s =========

Example

The examples/ folder contains an example for training an auto-encoder on MNIST.

Output of the example script. First two columns show training set, second two columns show testing set.

FAQ

1. Output is not equal to the input

Either normalization gone wrong or you encoded a symbol that is >Lp.

Important Implementation Details

How we represent probability distributions.

The probabilities are specified as CDFs. For each possible symbol, we need 2 CDF values. This means that if there are L possible symbols {0, ..., L-1}, the CDF must specified the value for L+1 symbols.

Example:

Let's say we have L = 3 possible symbols. We need a CDF with 4 values
to specify the symbols distribution:

symbol:        0     1     2
cdf:       C_0   C_1   C_2   C_3

This corresponds to the 3 probabilities

P(0) = C_1 - C_0
P(1) = C_2 - C_1
P(2) = C_3 - C_2

NOTE: The arithmetic coder assumes that C_3 == 1. 

Important:

• If you have L possible symbols, you need to pass a CDF that specifies L + 1 values. Since this is a common number, we call it Lp = L + 1 throught the code (the "p" stands for prime, i.e., L').
• The last value of the CDF should be 1. Note that the arithmetic coder in torchac.cpp will just assume it's 1 regardless of what is passed, so not having a CDF that ends in 1 will mean you will estimate bitrates wrongly. More details below.
• Note that even though the CDF specifies Lp values, symbols are only allowed to be in {0, ..., Lp-2}. In the above example, Lp == 4, but the max symbols is Lp-2 == 2. Bigger values will yield wrong outputs

Expected input shapes

We allow any shapes for the inputs, but the spatial dimensions of the input CDF and the input symbols must match. In particular, we expect:

• CDF must have shape (N1, ..., Nm, Lp), where N1, ..., Nm are the m spatial dimensions, and Lp is as described above.
• Symbols must have shape (N1, ..., Nm), i.e., same spatial dimensions as the CDF.

For example, in a typical CNN, you might have a CDF of shape (batch, channels, height, width, Lp).

Normalized vs. Unnormalized / Floating Point vs. Integer CDFs

The library differentiates between "normalized" and "unnormalized" CDFs, and between "floating point" and "integer" CDFs. What do these mean?

• A proper CDF is strictly monotonically increasing, and we call this a "normalized" CDF.
• However, since we work with finite precision (16 bits to be precise in this implementation), it may be that you have a CDF that is strictly monotonically increasing in float32 space, but not when it is converted to 16 bit precision. An "unnormalized" CDF is what we call a CDF that has the same value for at least two subsequent elements.
• "floating point" CDFs are CDFs that are specified as float32 and need to be converted to 16 bit precision
• "integer" CDFs are CDFs specified as int16 - BUT are then interpreted as uint16 on the C++ side. See "int16 vs uint16" below.

Examples:

float_unnormalized_cdf = [0.1, 0.2, 0.2, 0.3, ..., 1.]
float_normalized_cdf = [0.1, 0.2, 0.20001, 0.3, ..., 1.]
integer_unnormalized_cdf = [10, 20, 20, 30, ..., 0]  # See below for why last is 0.
integer_normalized_cdf = [10, 20, 21, 30, ..., 0]    # See below for why last is 0.

There are two APIs:

• encode_float_cdf and decode_float_cdf is to be used for floating point CDFs. These functions have a flag needs_normalization that specifies whether the input is assumed to be normalized. You can set need_normalization=False if you have CDFs that you know are normalized, e.g., Gaussian distributions with a large enough sigma. This would then speedup encoding and decoding large tensors somewhat, and will make bitrate estimation from the CDF more precise.
• encode_int16_normalized_cdf and decode_int16_normalized_cdf is to be used for integer CDFs that are already normalized.

int16 vs uint16 - it gets confusing!

One big source of confusion can be that PyTorch does not support uint16. Yet, that's exactly what we need. So what we do is we just represent integer CDFs with int16 in the Python side, and interpret/cast them to uint16 on the C++ side. This means that if you were to look at the int16 CDFs you would see confusing things:

# Python
cdf_float = [0., 1/3, 2/3, 1.]  # A uniform distribution for L=3 symbols.
cdf_int = [0, 21845, -21845, 0]

# C++
uint16* cdf_int = [0, 21845, 43690, 0]

Note:

1. In the python cdf_int numbers bigger than 2**16/2 are negative
2. The final value is actually 0. This is then handled in torchac.cpp which just assums cdf[..., -1] == 2**16, which cannot be represented as a uint16.

Fun stuff!