Package: torchac
torchac: Fast Arithmetic Coding for PyTorchAboutThis is a stand-alone version of the arithmetic coder we used in the
neural compression paperPractical Full Resolution Learned Lossless Image
Compressionby Mentzer et al.The backend is written in C++, the API is for PyTorch tensors. Thanks to
on-the-fly compilation with ninja, the integration is seamless.The implementation is based onthis blog post,
meaning that we implementarithmetic 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 all pixels of a512 x 512image happens in 0.202s (seeAppendix Ain the paper).What torchac isA simple-to-use library to encode a stream of symbols into a bitstream given
the cumulative distribution (CDF) of the symbols.
The number of possible symbols must be finite.What torchac is notWe do not provide classes to learn or represent probability/cumulative
distributions. These have to be provided by you.HowToSet up conda environmentThis library has been tested withPyTorch 1.5, 1.6, 1.7Python 3.8And that's all you need. Other versions of Python may also work,
but on-the-fly ninja compilation only works for PyTorch 1.5+.In a supported environment, installtorchacwithpip:pipinstalltorchacIf you don't have an environment already set up, you can make one withconda,
seepytorch.org.To (optionally) test your installation, you needpytest, and then you
can runpython -m pytest tests/test.py -s. The output should contain5 passed, 2 warnings.ExampleTheexamples/folder containsan example for training an auto-encoder on MNIST.Snipped from that example:importtorchac# Encode to bytestream.output_cdf=...# Get CDF from your model, shape B, C, H, W, Lpsym=...# Get the symbols to encode, shape B, C, H, W.byte_stream=torchac.encode_float_cdf(output_cdf,sym,check_input_bounds=True)# Number of bits taken by the streamreal_bits=len(byte_stream)*8# Write to a file.withopen('outfile.b','wb')asfout:fout.write(byte_stream)# Read from a file.withopen('outfile.b','rb')asfin:byte_stream=fin.read()# Decode from bytestream.sym_out=torchac.decode_float_cdf(output_cdf,byte_stream)# Output will be equal to the input.assertsym_out.equal(sym)FAQ1. Output is not equal to the inputEither normalization went wrong or you encoded a symbol that is>Lp,
see below for more details.Important Implementation DetailsHow we represent probability distributions.The probabilities are specified asCDFs.
For each possible symbol,
we need 2 CDF values. This means that if there areLpossible symbols{0, ..., L-1}, the CDF must specified the value forL+1symbols.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 haveLpossible symbols, you need to pass a CDF that
specifiesL + 1values. Since this is a common number, we call itLp = L + 1throught the code (the "p" stands for prime, i.e.,L').The last value of the CDF should be1. Note that the arithmetic coder
intorchac.cppwill just assume it's1regardless of what is passed, so not having a CDF
that ends in1will mean you will estimate bitrates wrongly. More details below.Note that even though the CDF specifiesLpvalues, symbols are only allowed
to be in{0, ..., Lp-2}. In the above example,Lp == 4, but the
max symbols isLp-2 == 2. Bigger values will yieldwrong outputsExpected input shapesWe 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), whereN1, ..., Nmare themspatial dimensions, andLpis 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 CDFsThe 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 infloat32space, 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 asfloat32and need
to be converted to 16 bit precision."integer" CDFs are CDFs specified asint16- BUT are then interpreted
asuint16on 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_cdfanddecode_float_cdfis to be used for floating point
CDFs. These functions have a flagneeds_normalizationthat specifies
whether the input is assumed to be normalized. You can setneed_normalization=Falseif 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_cdfanddecode_int16_normalized_cdfis to be
used for integer CDFsthat are already normalized.int16 vs uint16 - it gets confusing!One big source of confusion can be that PyTorch does not supportuint16.
Yet, that's exactly what we need. So what we do is we just represent
integer CDFs withint16in the Python side, and interpret/cast them touint16on the C++ side. This means that if you were to look at the int16 CDFs
you would see confusing things:# Pythoncdf_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:In the pythoncdf_intnumbers bigger than2**16/2are negativeThe final value is actually 0. This is then handled intorchac.cppwhich
just assumscdf[..., -1] == 2**16, which cannot be represented as auint16.Fun stuff!CitationIf you use the work released here for your research, consider citing this paper:@inproceedings{mentzer2019practical,
    Author = {Mentzer, Fabian and Agustsson, Eirikur and Tschannen, Michael and Timofte, Radu and Van Gool, Luc},
    Booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    Title = {Practical Full Resolution Learned Lossless Image Compression},
    Year = {2019}}