2 | 
3 |
4 |
5 |
6 | # torchac: Fast Arithmetic Coding for PyTorch
7 |
8 | [](https://circleci.com/gh/fab-jul/torchac)
9 | [](https://pypi.org/project/torchac/)
10 | 
11 | 
12 |
13 |
14 |
15 |
16 | ## About
17 |
18 | This is a stand-alone version of the arithmetic coder we used in the
19 | neural compression paper [Practical Full Resolution Learned Lossless Image
20 | Compression](https://github.com/fab-jul/L3C-PyTorch) by Mentzer et al.
21 |
22 | The backend is written in C++, the API is for PyTorch tensors. Thanks to
23 | on-the-fly compilation with ninja, the integration is seamless.
24 |
25 | The implementation is based on [this blog post](https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html),
26 | meaning that we implement _arithmetic coding_.
27 | While it could be further optimized, it is already much faster than doing the equivalent thing in pure-Python (because of all the
28 | bit-shifts etc.). In L3C, Encoding all pixels of a `512 x 512` image happens in 0.202s (see [Appendix A](https://arxiv.org/abs/1811.12817) in the paper).
29 |
30 | ### What torchac is
31 |
32 | - A simple-to-use library to encode a stream of symbols into a bitstream given
33 | the cumulative distribution (CDF) of the symbols.
34 | The number of possible symbols must be finite.
35 |
36 | ### What torchac is not
37 |
38 | - We do not provide classes to learn or represent probability/cumulative
39 | distributions. These have to be provided by you.
40 |
41 |
42 | ## HowTo
43 |
44 | ### Set up conda environment
45 |
46 | This library has been tested with
47 | - PyTorch 1.5 - 1.12
48 | - Python 3.8, 3.9
49 |
50 | Other versions of Python may also work,
51 | but on-the-fly ninja compilation only works for PyTorch 1.5+.
52 |
53 | In a supported environment, install `torchac` with `pip`:
54 |
55 | ```bash
56 | pip install torchac
57 | ```
58 |
59 | If you don't have an environment already set up, you can make one with `conda`,
60 | see [pytorch.org](https://pytorch.org).
61 |
62 | #### Testing installation
63 |
64 | To test the installation, `git clone` this repo and run `bash
65 | install_and_run_test.sh`. It should end in a line that says that `5 passed`.
66 |
67 | ### Example
68 |
69 | The `examples/` folder contains [an example for training an auto-encoder on MNIST](https://github.com/fab-jul/torchac/tree/master/examples/mnist_autoencoder).
70 |
71 |
3 |
72 | 
73 |
76 |
77 | Snipped from that example:
78 |
79 | ```python
80 | import torchac
81 |
82 | # Encode to bytestream.
83 | output_cdf = ... # Get CDF from your model, shape B, C, H, W, Lp
84 | sym = ... # Get the symbols to encode, shape B, C, H, W.
85 | byte_stream = torchac.encode_float_cdf(output_cdf, sym, check_input_bounds=True)
86 |
87 | # Number of bits taken by the stream
88 | real_bits = len(byte_stream) * 8
89 |
90 | # Write to a file.
91 | with open('outfile.b', 'wb') as fout:
92 | fout.write(byte_stream)
93 |
94 | # Read from a file.
95 | with open('outfile.b', 'rb') as fin:
96 | byte_stream = fin.read()
97 |
98 | # Decode from bytestream.
99 | sym_out = torchac.decode_float_cdf(output_cdf, byte_stream)
100 |
101 | # Output will be equal to the input.
102 | assert sym_out.equal(sym)
103 | ```
104 |
105 | ## FAQ
106 |
107 | #### 1. Output is not equal to the input
108 |
109 | Either normalization went wrong or you encoded a symbol that is `>Lp`,
110 | see below for more details.
111 |
112 | ## Important Implementation Details
113 |
114 | ### How we represent probability distributions.
115 |
116 | The probabilities are specified as [CDFs](https://en.wikipedia.org/wiki/Cumulative_distribution_function).
117 | For each possible symbol,
118 | we need 2 CDF values. This means that if there are `L` possible symbols
119 | `{0, ..., L-1}`, the CDF must specified the value for `L+1` symbols.
120 |
121 | **Example**:
122 | ```
123 | Let's say we have L = 3 possible symbols. We need a CDF with 4 values
124 | to specify the symbols distribution:
125 |
126 | symbol: 0 1 2
127 | cdf: C_0 C_1 C_2 C_3
128 |
129 | This corresponds to the 3 probabilities
130 |
131 | P(0) = C_1 - C_0
132 | P(1) = C_2 - C_1
133 | P(2) = C_3 - C_2
134 |
135 | NOTE: The arithmetic coder assumes that C_3 == 1.
136 | ```
137 |
138 | Important:
139 |
140 | - If you have `L` possible symbols, you need to pass a CDF that
141 | specifies `L + 1` values. Since this is a common number, we call it
142 | `Lp = L + 1` throught the code (the "p" stands for prime, i.e., `L'`).
143 | - The last value of the CDF should be `1`. Note that the arithmetic coder
144 | in `torchac.cpp` will just assume it's `1` regardless of what is passed, so not having a CDF
145 | that ends in `1` will mean you will estimate bitrates wrongly. More details below.
146 | - Note that even though the CDF specifies `Lp` values, symbols are only allowed
147 | to be in `{0, ..., Lp-2}`. In the above example, `Lp == 4`, but the
148 | max symbols is `Lp-2 == 2`. Bigger values will yield **wrong outputs**
149 |
150 | ### Expected input shapes
151 |
152 | We allow any shapes for the inputs, but the spatial dimensions of the
153 | input CDF and the input symbols must match. In particular, we expect:
154 |
155 | - CDF must have shape `(N1, ..., Nm, Lp)`, where `N1, ..., Nm` are the
156 | `m` spatial dimensions, and `Lp` is as described above.
157 | - Symbols must have shape `(N1, ..., Nm)`, i.e., same spatial dimensions
158 | as the CDF.
159 |
160 | For example, in a typical CNN, you might have a CDF of shape
161 | `(batch, channels, height, width, Lp)`.
162 |
163 |
164 | ### Normalized vs. Unnormalized / Floating Point vs. Integer CDFs
165 |
166 | The library differentiates between "normalized" and "unnormalized" CDFs,
167 | and between "floating point" and "integer" CDFs. What do these mean?
168 |
169 | - A proper CDF is strictly monotonically increasing, and we call this a
170 | "normalized" CDF.
171 | - However, since we work with finite precision (16 bits to
172 | be precise in this implementation), it may be that you have a CDF that
173 | is strictly monotonically increasing in `float32` space, but not when
174 | it is converted to 16 bit precision. An "unnormalized" CDF is what we call
175 | a CDF that has the same value for at least two subsequent elements.
176 | - "floating point" CDFs are CDFs that are specified as `float32` and need
177 | to be converted to 16 bit precision.
178 | - "integer" CDFs are CDFs specified as `int16` - BUT are then interpreted
179 | as `uint16` on the C++ side. See "int16 vs uint16" below.
180 |
181 | Examples:
182 |
183 | ```python
184 | float_unnormalized_cdf = [0.1, 0.2, 0.2, 0.3, ..., 1.]
185 | float_normalized_cdf = [0.1, 0.2, 0.20001, 0.3, ..., 1.]
186 | integer_unnormalized_cdf = [10, 20, 20, 30, ..., 0] # See below for why last is 0.
187 | integer_normalized_cdf = [10, 20, 21, 30, ..., 0] # See below for why last is 0.
188 | ```
189 |
190 | There are two APIs:
191 |
192 | - `encode_float_cdf` and `decode_float_cdf` is to be used for floating point
193 | CDFs. These functions have a flag `needs_normalization` that specifies
194 | whether the input is assumed to be normalized. You can set
195 | `need_normalization=False` if you have CDFs that you know are normalized, e.g.,
196 | Gaussian distributions with a large enough sigma. This would then speedup
197 | encoding and decoding large tensors somewhat, and will make bitrate
198 | estimation from the CDF more precise.
199 | - `encode_int16_normalized_cdf` and `decode_int16_normalized_cdf` is to be
200 | used for integer CDFs **that are already normalized**.
201 |
202 | ### int16 vs uint16 - it gets confusing!
203 |
204 | One big source of confusion can be that PyTorch does not support `uint16`.
205 | Yet, that's exactly what we need. So what we do is we just represent
206 | integer CDFs with `int16` in the Python side, and interpret/cast them to `uint16`
207 | on the C++ side. This means that if you were to look at the int16 CDFs
208 | you would see confusing things:
209 |
210 | ```python
211 | # Python
212 | cdf_float = [0., 1/3, 2/3, 1.] # A uniform distribution for L=3 symbols.
213 | cdf_int = [0, 21845, -21845, 0]
214 |
215 | # C++
216 | uint16* cdf_int = [0, 21845, 43690, 0]
217 | ```
218 |
219 | Note:
220 | 1. In the python `cdf_int` numbers bigger than `2**16/2` are negative
221 | 2. The final value is actually 0. This is then handled in `torchac.cpp` which
222 | just assums `cdf[..., -1] == 2**16`, which cannot be represented as a `uint16`.
223 |
224 | Fun stuff!
225 |
226 | ## Citation
227 |
228 | If you use the work released here for your research, consider citing this paper:
229 | ```
230 | @inproceedings{mentzer2019practical,
231 | Author = {Mentzer, Fabian and Agustsson, Eirikur and Tschannen, Michael and Timofte, Radu and Van Gool, Luc},
232 | Booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
233 | Title = {Practical Full Resolution Learned Lossless Image Compression},
234 | Year = {2019}}
235 | ```
236 |
237 |
--------------------------------------------------------------------------------
/TODO.md:
--------------------------------------------------------------------------------
1 | ## **State**: Released
2 | #### Future
3 |
4 | - [ ] Docker container with all dependencies
5 | - [ ] link to here from L3C
6 | - [ ] estimate bitrate from normalized CDF
7 |
--------------------------------------------------------------------------------
/bin/deploy.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | set -e
4 |
5 | VERSION_NUMBER=$1
6 |
7 | if [[ -z $VERSION_NUMBER ]]; then
8 | echo "Usage: $0 VERSION_NUMBER"
9 | exit 1
10 | fi
11 |
12 | if [ -n "$(git status --porcelain)" ]; then
13 | echo "Error: Git not clean, please commit."
14 | exit 1
15 | fi
16 |
17 | python bin/update_version.py $VERSION_NUMBER
18 |
19 | rm -rf dist/
20 | python setup.py bdist_wheel
21 | twine upload dist/*
22 |
23 | python bin/update_version.py $VERSION_NUMBER --set-used
24 |
25 | echo "Waiting for changes to propagate..."
26 | sleep 60 # Give PyPi time to prepare the package.
27 |
28 | bash pypi/test.sh tests/test.py
29 |
30 | git commit bin/version.json -m "Updated Version: $VERSION_NUMBER"
31 | git tag $VERSION_NUMBER
32 | git push
33 | git push --tags
34 |
--------------------------------------------------------------------------------
/bin/update_version.py:
--------------------------------------------------------------------------------
1 | from distutils.version import LooseVersion, StrictVersion
2 | import argparse
3 | import json
4 |
5 |
6 | def main():
7 | p = argparse.ArgumentParser()
8 | p.add_argument('new_version')
9 | p.add_argument('--set-used', action='store_true')
10 | flags = p.parse_args()
11 | new_version = flags.new_version
12 | set_used = flags.set_used
13 |
14 | with open('bin/version.json', 'r') as f:
15 | version_info = json.load(f)
16 |
17 | is_unused = not version_info["used"]
18 | cur_version = version_info["version"]
19 | if is_unused:
20 | version_is_ok = LooseVersion(new_version) >= LooseVersion(cur_version)
21 | else:
22 | version_is_ok = LooseVersion(new_version) > LooseVersion(cur_version)
23 |
24 | if not version_is_ok:
25 | raise ValueError(f'Not ok: {new_version}. Current version={cur_version}')
26 |
27 | with open('bin/version.json', 'w') as f:
28 | json.dump({'version': new_version, 'used': set_used}, f)
29 |
30 |
31 | if __name__ == '__main__':
32 | main()
33 |
34 |
35 |
36 |
--------------------------------------------------------------------------------
/bin/version.json:
--------------------------------------------------------------------------------
1 | {"version": "0.9.3", "used": true}
--------------------------------------------------------------------------------
/examples/mnist_autoencoder/README.md:
--------------------------------------------------------------------------------
1 | # MNIST AutoEncoder
2 |
3 |
73 | Output of the example script. First two columns show training 74 | set, second two columns show testing set.
75 |
4 | 
5 |
8 |
9 | This example shows how to train a simple auto-encoder for MNIST,
10 | with a quantized bottleneck, and a conditional probability
11 | model that estimates the distribution of the bottleneck. It basically
12 | follows recent image compression papers in terms of structure,
13 | but is much simpler and not tuned.
14 |
15 | Goals:
16 | - Show how to use the probability distribution predicted from some probability
17 | model CNN to estimate the bitrate
18 | - Actually encode symbols with that model using `torchac`, and see how many
19 | bits it takes to store
20 |
21 | What could be improved:
22 | - The models don't actually train very well. The purpose of the example is to
23 | show how `torchac` can be used.
24 |
25 | ## Runnig it
26 |
27 | The example supports interactive plots with matplotlib, which it tries
28 | to use by default. If they are not available (in headless environments),
29 | it falls back to storing plots in `plots`. If you want to force
30 | non-interactive, you can set `export NO_INTERACTIVE=1` before
31 | running the script.
32 |
33 | ```bash
34 | # If not installed already:
35 | pip install matplotlib
36 | pip install torchac
37 |
38 | # Got to this folder
39 | cd examples/mnist_qutoencder
40 |
41 | # Run the script
42 | python mnist_autoencoder_example.py
43 | ```
44 |
--------------------------------------------------------------------------------
/examples/mnist_autoencoder/mnist_autoencoder_example.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import dataclasses
3 | import itertools
4 | import warnings
5 |
6 | import os
7 | import time
8 |
9 | import torch
10 | from torch import nn
11 | import torch.nn.functional as F
12 |
13 | import numpy as np
14 | import torchvision
15 | import matplotlib.pyplot as plt
16 | import matplotlib
17 |
18 | try:
19 | import torchac
20 | except ImportError:
21 | raise ImportError('torchac is not available! Please see the main README '
22 | 'on how to install it.')
23 |
24 |
25 | # Set to true to write to disk.
26 | _WRITE_BITS = False
27 |
28 |
29 | # Interactive plots setup
30 | _FORCE_NO_INTERACTIVE_PLOTS = int(os.environ.get('NO_INTERACTIVE', 0)) == 1
31 |
32 |
33 | if not _FORCE_NO_INTERACTIVE_PLOTS:
34 | try:
35 | matplotlib.use("TkAgg")
36 | interactive_plots_available = True
37 | except ImportError:
38 | warnings.warn(f'*** TkAgg not available! Saving plots...')
39 | interactive_plots_available = False
40 | else:
41 | interactive_plots_available = False
42 |
43 | if not interactive_plots_available:
44 | matplotlib.use("Agg")
45 |
46 |
47 | # Set seed.
48 | torch.manual_seed(0)
49 |
50 |
51 | def train_test_loop(use_gpu=True,
52 | bottleneck_size=2,
53 | L=5,
54 | batch_size=8,
55 | lr=1e-2,
56 | rate_loss_enable_itr=500,
57 | num_test_batches=10,
58 | train_plot_every_itr=50,
59 | max_training_itr=None,
60 | mnist_download_dir='data',
61 | ):
62 | """Train and test an autoencoder.
63 |
64 | :param use_gpu: Whether to use the GPU, if it is available.
65 | :param bottleneck_size: Number of channels in the bottleneck.
66 | :param L: Number of levels that we quantize to.
67 | :param batch_size: Batch size we train with.
68 | :param lr: Learning rate of Adam.
69 | :param rate_loss_enable_itr: Iteration when the rate loss is enabled.
70 | :param num_test_batches: Number of batches we test on (randomly chosen).
71 | :param train_plot_every_itr: How often to update the train plot.
72 | :param max_training_itr: If given, only train for max_training_itr iterations.
73 | :param mnist_download_dir: Where to store MNIST.
74 | """
75 | ae = Autoencoder(bottleneck_size, L)
76 | prob = ConditionalProbabilityModel(L=L, bottleneck_shape=ae.bottleneck_shape)
77 |
78 | device = 'cuda' if (use_gpu and torch.cuda.is_available()) else 'cpu'
79 | ae = ae.to(device)
80 | prob = prob.to(device)
81 |
82 | mse = nn.MSELoss()
83 | adam = torch.optim.Adam(
84 | itertools.chain(ae.parameters(), prob.parameters()),
85 | lr=lr)
86 |
87 | train_acc = Accumulator()
88 | test_acc = Accumulator()
89 | plotter = Plotter()
90 |
91 | transforms = torchvision.transforms.Compose([
92 | torchvision.transforms.ToTensor(),
93 | torchvision.transforms.Lambda(
94 | # Make images 32x32.
95 | lambda image: F.pad(image, pad=(2, 2, 2, 2), mode='constant'))
96 | ])
97 |
98 | train_loader = torch.utils.data.DataLoader(
99 | torchvision.datasets.MNIST(mnist_download_dir, train=True, download=True,
100 | transform=transforms),
101 | batch_size=batch_size, shuffle=True)
102 |
103 | rate_loss_enabled = False
104 | for i, (images, labels) in enumerate(train_loader):
105 | if max_training_itr and i >= max_training_itr:
106 | break
107 | assert images.shape[-2:] == (32, 32)
108 | images = images.to(device)
109 | labels = labels.to(device)
110 |
111 | adam.zero_grad()
112 |
113 | # Get reconstructions and symbols from the autoencoder.
114 | reconstructions, sym = ae(images)
115 | assert sym.shape[1:] == ae.bottleneck_shape
116 |
117 | # Get estimated and real bitrate from probability model, given labels.
118 | bits_estimated, bits_real = prob(sym.detach(), labels)
119 | mse_loss = mse(reconstructions, images)
120 |
121 | # If we are beyond iteration `rate_loss_enable_itr`, enable a rate loss.
122 | if i < rate_loss_enable_itr:
123 | loss = mse_loss
124 | else:
125 | loss = mse_loss + 1/1000 * bits_estimated
126 | rate_loss_enabled = True
127 |
128 | loss.backward()
129 | adam.step()
130 |
131 | # Update Train Plot.
132 | if i > 0 and i % train_plot_every_itr == 0:
133 | train_acc.append(i, bits_estimated, bits_real, mse_loss)
134 | print(f'{i: 10d}; '
135 | f'loss={loss:.3f}, '
136 | f'bits_estimated={bits_estimated:.3f}, '
137 | f'mse={mse_loss:.3f}')
138 | plotter.update('Train',
139 | images, reconstructions, sym, train_acc, rate_loss_enabled)
140 |
141 | # Update Test Plot.
142 | if i > 0 and i % 100 == 0:
143 | print(f'{i: 10d} Testing on {num_test_batches} random batches...')
144 | ae.eval()
145 | prob.eval()
146 | test_loader = torch.utils.data.DataLoader(
147 | torchvision.datasets.MNIST(
148 | mnist_download_dir, train=False, download=True, transform=transforms),
149 | batch_size=batch_size, shuffle=True)
150 | with torch.no_grad():
151 | across_batch_acc = Accumulator()
152 | for j, (test_images, test_labels) in enumerate(test_loader):
153 | if j >= num_test_batches:
154 | break
155 | test_images = test_images.to(device)
156 | test_labels = test_labels.to(device)
157 | test_reconstructions, test_sym = ae(test_images)
158 | test_bits_estimated, test_bits_real = prob(test_sym, test_labels)
159 | test_mse_loss = mse(test_reconstructions, test_images)
160 | across_batch_acc.append(j, test_bits_estimated, test_bits_real, test_mse_loss)
161 | test_bits_estimated_mean, test_bits_real_mean, test_mse_loss_mean = \
162 | across_batch_acc.means()
163 | test_acc.append(
164 | i, test_bits_estimated_mean, test_bits_real_mean, test_mse_loss_mean)
165 | plotter.update('Test', test_images, test_reconstructions, test_sym,
166 | test_acc, rate_loss_enabled)
167 | ae.train()
168 | prob.train()
169 |
170 |
171 | class STEQuantize(torch.autograd.Function):
172 | """Straight-Through Estimator for Quantization.
173 |
174 | Forward pass implements quantization by rounding to integers,
175 | backward pass is set to gradients of the identity function.
176 | """
177 | @staticmethod
178 | def forward(ctx, x):
179 | ctx.save_for_backward(x)
180 | return x.round()
181 |
182 | @staticmethod
183 | def backward(ctx, grad_outputs):
184 | return grad_outputs
185 |
186 |
187 | class Autoencoder(nn.Module):
188 | def __init__(self, bottleneck_size, L):
189 | if L % 2 != 1:
190 | raise ValueError(f'number of levels L={L}, must be odd number!')
191 | super(Autoencoder, self).__init__()
192 | self.L = L
193 | self.enc = nn.Sequential(
194 | nn.Conv2d(1, 32, 5, stride=2, padding=2),
195 | nn.InstanceNorm2d(32),
196 | nn.ReLU(),
197 | nn.Conv2d(32, 32, 5, stride=2, padding=2),
198 | nn.InstanceNorm2d(32),
199 | nn.ReLU(),
200 | nn.Conv2d(32, 32, 5, stride=2, padding=2),
201 | nn.InstanceNorm2d(32),
202 | nn.ReLU(),
203 | nn.Conv2d(32, 32, 5, stride=2, padding=2),
204 | nn.InstanceNorm2d(32),
205 | nn.ReLU(),
206 | nn.Conv2d(32, bottleneck_size, 1, stride=1, padding=0, bias=False),
207 | )
208 |
209 | self.dec = nn.Sequential(
210 | nn.ConvTranspose2d(bottleneck_size, 32, 5, stride=2, padding=2, output_padding=1),
211 | nn.InstanceNorm2d(32),
212 | nn.ReLU(),
213 | nn.ConvTranspose2d(32, 32, 5, stride=2, padding=2, output_padding=1),
214 | nn.InstanceNorm2d(32),
215 | nn.ReLU(),
216 | nn.ConvTranspose2d(32, 32, 5, stride=2, padding=2, output_padding=1),
217 | nn.InstanceNorm2d(32),
218 | nn.ReLU(),
219 | nn.ConvTranspose2d(32, 32, 5, stride=2, padding=2, output_padding=1),
220 | nn.InstanceNorm2d(32),
221 | nn.ReLU(),
222 |
223 | # Add a few convolutions at the final resolution.
224 | nn.Conv2d(32, 32, 3, stride=1, padding=1),
225 | nn.ReLU(),
226 | nn.Conv2d(32, 32, 1, stride=1, padding=0),
227 | nn.ReLU(),
228 | nn.Conv2d(32, 1, 1, stride=1),
229 | )
230 |
231 | self.quantize = STEQuantize.apply
232 | self.bottleneck_shape = (bottleneck_size, 2, 2)
233 |
234 | def forward(self, image):
235 | # Encode image x into the latent.
236 | latent = self.enc(image)
237 | # The jiggle is there so that the lowest and highest symbol are not at
238 | # the boundary. Probably not needed.
239 | jiggle = 0.2
240 | spread = self.L - 1 + jiggle
241 | # The sigmoid clamps to [0, 1], then we multiply it by spread and substract
242 | # spread / 2, so that the output is nicely centered around zero and
243 | # in the interval [-spread/2, spread/2]
244 | latent = torch.sigmoid(latent) * spread - spread / 2
245 | latent_quantized = self.quantize(latent)
246 | reconstructions = self.dec(latent_quantized)
247 | sym = latent_quantized + self.L // 2
248 | return reconstructions, sym.to(torch.long)
249 |
250 |
251 | class ConditionalProbabilityModel(nn.Module):
252 | def __init__(self, L, bottleneck_shape):
253 | super(ConditionalProbabilityModel, self).__init__()
254 | self.L = L
255 | self.bottleneck_shape = bottleneck_shape
256 |
257 | self.bottleneck_size, _, _ = bottleneck_shape
258 |
259 | # We predict a value for each channel, for each level.
260 | num_output_channels = self.bottleneck_size * L
261 |
262 | self.model = nn.Sequential(
263 | nn.Conv2d(1, self.bottleneck_size, 3, padding=1),
264 | nn.BatchNorm2d(self.bottleneck_size),
265 | nn.ReLU(),
266 | nn.Conv2d(self.bottleneck_size, self.bottleneck_size, 3, padding=1),
267 | nn.BatchNorm2d(self.bottleneck_size),
268 | nn.ReLU(),
269 | nn.Conv2d(self.bottleneck_size, num_output_channels, 1, padding=0),
270 | )
271 |
272 | def forward(self, sym, labels):
273 | batch_size = sym.shape[0]
274 | _, H, W = self.bottleneck_shape
275 | # Construct the input, which is just the label of the current number
276 | # at each spatial location.
277 | bottleneck_shape_with_batch_dim = (batch_size, 1, H, W)
278 | static_input = torch.ones(
279 | bottleneck_shape_with_batch_dim, dtype=torch.float32, device=sym.device)
280 | dynamic_input = static_input * labels.reshape(-1, 1, 1, 1)
281 | # Divide by 9 and substract 0.5 to center the input around 0 and make
282 | # it be contained in [-0.5, 0.5].
283 | dynamic_input = dynamic_input / 9 - 0.5
284 |
285 | # Get the output of the CNN.
286 | output = self.model(dynamic_input)
287 | _, C, H, W = output.shape
288 | assert C == self.bottleneck_size * self.L
289 |
290 | # Reshape it such that the probability per symbol has it's own dimension.
291 | # output_reshaped has shape (B, C, L, H, W).
292 | output_reshaped = output.reshape(
293 | batch_size, self.bottleneck_size, self.L, H, W)
294 | # Take the softmax over that dimension to make this into a normalized
295 | # probability distribution.
296 | output_probabilities = F.softmax(output_reshaped, dim=2)
297 | # Permute the symbols dimension to the end, as expected by torchac.
298 | # output_probabilities has shape (B, C, H, W, L).
299 | output_probabilities = output_probabilities.permute(0, 1, 3, 4, 2)
300 | # Estimate the bitrate from the PMF.
301 | estimated_bits = estimate_bitrate_from_pmf(output_probabilities, sym=sym)
302 | # Convert to a torchac-compatible CDF.
303 | output_cdf = pmf_to_cdf(output_probabilities)
304 | # torchac expects sym as int16, see README for details.
305 | sym = sym.to(torch.int16)
306 | # torchac expects CDF and sym on CPU.
307 | output_cdf = output_cdf.detach().cpu()
308 | sym = sym.detach().cpu()
309 | # Get real bitrate from the byte_stream.
310 | byte_stream = torchac.encode_float_cdf(output_cdf, sym, check_input_bounds=True)
311 | real_bits = len(byte_stream) * 8
312 | if _WRITE_BITS:
313 | # Write to a file.
314 | with open('outfile.b', 'wb') as fout:
315 | fout.write(byte_stream)
316 | # Read from a file.
317 | with open('outfile.b', 'rb') as fin:
318 | byte_stream = fin.read()
319 | assert torchac.decode_float_cdf(output_cdf, byte_stream).equal(sym)
320 | return estimated_bits, real_bits
321 |
322 |
323 | def pmf_to_cdf(pmf):
324 | cdf = pmf.cumsum(dim=-1)
325 | spatial_dimensions = pmf.shape[:-1] + (1,)
326 | zeros = torch.zeros(spatial_dimensions, dtype=pmf.dtype, device=pmf.device)
327 | cdf_with_0 = torch.cat([zeros, cdf], dim=-1)
328 | # On GPU, softmax followed by cumsum can lead to the final value being
329 | # slightly bigger than 1, so we clamp.
330 | cdf_with_0 = cdf_with_0.clamp(max=1.)
331 | return cdf_with_0
332 |
333 |
334 | def estimate_bitrate_from_pmf(pmf, sym):
335 | L = pmf.shape[-1]
336 | pmf = pmf.reshape(-1, L)
337 | sym = sym.reshape(-1, 1)
338 | assert pmf.shape[0] == sym.shape[0]
339 | relevant_probabilities = torch.gather(pmf, dim=1, index=sym)
340 | bitrate = torch.sum(-torch.log2(relevant_probabilities.clamp(min=1e-3)))
341 | return bitrate
342 |
343 |
344 | @dataclasses.dataclass
345 | class Accumulator:
346 | def __init__(self):
347 | self.iterations = []
348 | self.bits_estimated_acc = []
349 | self.bits_real_acc = []
350 | self.mse_acc = []
351 |
352 | def append(self, i, bits_estimated, bits_real, mse):
353 | self.iterations.append(i)
354 | self.bits_estimated_acc.append(bits_estimated.item())
355 | self.bits_real_acc.append(bits_real)
356 | self.mse_acc.append(mse.item())
357 |
358 | def get_errors(self):
359 | return [real / estimated - 1 for real, estimated in
360 | zip(self.bits_real_acc, self.bits_estimated_acc)]
361 |
362 | def means(self):
363 | return (np.mean(self.bits_estimated_acc),
364 | np.mean(self.bits_real_acc),
365 | np.mean(self.mse_acc))
366 |
367 |
368 | class Plotter(object):
369 | def __init__(self):
370 | plt.ion()
371 | self.fig, axs = plt.subplots(ncols=4, nrows=4, figsize=(12, 8))
372 | Plotter._setup_axes(axs[:, :2], 'Train')
373 | Plotter._setup_axes(axs[:, 2:], 'Test')
374 | plt.tight_layout()
375 | if interactive_plots_available:
376 | plt.draw()
377 | else:
378 | unique_id = str(time.time()).replace('.', '_')
379 | self.out_dir = os.path.join('plots', unique_id)
380 | os.makedirs(self.out_dir, exist_ok=True)
381 | self.axs = axs
382 |
383 | @staticmethod
384 | def _setup_axes(axs, title):
385 | axs[0, 0].set_axis_off()
386 | axs[0, 0].set_title(f'{title} Input')
387 | axs[0, 1].set_axis_off()
388 | axs[0, 1].set_title(f'{title} Reconstruction')
389 | axs[1, 0].set_axis_off()
390 | axs[1, 0].set_title(f'{title} Bottleneck channel 1')
391 | axs[1, 1].set_axis_off()
392 | axs[1, 1].set_title(f'{title} Bottleneck channel 2')
393 | axs[2, 0].set_title(f'Estimated Bitrate {title}')
394 | axs[2, 1].set_title(f'Real Bitrate {title}')
395 | axs[3, 0].set_title(f'MSE {title}')
396 | axs[3, 1].set_title('Rel. Bitrate Error')
397 |
398 | def update(self,
399 | mode,
400 | images, reconstructions, sym,
401 | acc: Accumulator,
402 | rate_loss_enabled: bool):
403 | # First two columns are for training, second two for testing.
404 | axs = self.axs[:, :2] if mode == 'Train' else self.axs[:, 2:]
405 | title = mode
406 |
407 | # Plot images.
408 | axs[0, 0].imshow(images[0, ...].squeeze().detach().cpu().numpy())
409 | axs[0, 1].imshow(reconstructions[0, ...].squeeze().detach().cpu().numpy())
410 | axs[1, 0].imshow(sym[0, 0, ...].to(torch.float).detach().cpu().numpy())
411 | axs[1, 1].imshow(sym[0, 1, ...].to(torch.float).detach().cpu().numpy())
412 |
413 | # Plot lines, make sure to empty plot first.
414 | axs[2, 0].clear()
415 | axs[2, 1].clear()
416 | axs[3, 0].clear()
417 | axs[3, 1].clear()
418 | axs[2, 0].set_title(f'Estimated Bitrate {title}')
419 | axs[2, 1].set_title(f'Real Bitrate {title}')
420 | axs[3, 0].set_title(f'MSE {title}')
421 | axs[3, 1].set_title('Rel. Bitrate Error')
422 |
423 | linestyle = '-' if rate_loss_enabled else ':'
424 | color = 'b' if mode == 'Train' else 'r'
425 | linestyle = color + linestyle
426 | axs[2, 0].plot(
427 | acc.iterations, acc.bits_estimated_acc, linestyle)
428 | axs[2, 1].plot(
429 | acc.iterations, acc.bits_real_acc, linestyle)
430 |
431 | axs[3, 0].plot(
432 | acc.iterations, acc.mse_acc, color)
433 | axs[3, 1].plot(
434 | acc.iterations, acc.get_errors(), color)
435 |
436 | if interactive_plots_available:
437 | plt.pause(0.05)
438 | else:
439 | plotname = f'{acc.iterations[-1]:010d}.png'
440 | out_p = os.path.join(self.out_dir, plotname)
441 | print(f'Saving plot at {out_p}...')
442 | self.fig.savefig(out_p)
443 |
444 |
445 | def main():
446 | p = argparse.ArgumentParser()
447 | p.add_argument('--max_training_itr', type=int)
448 | flags = p.parse_args()
449 | train_test_loop(max_training_itr=flags.max_training_itr)
450 |
451 |
452 | if __name__ == '__main__':
453 | main()
454 |
--------------------------------------------------------------------------------
/examples/mnist_autoencoder/progress_plot.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fab-jul/torchac/63064ab4caeaa7d435d7dc6877160161066cb558/examples/mnist_autoencoder/progress_plot.png
--------------------------------------------------------------------------------
/install_and_run_test.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | TEST_DIR=_torchac_test_
4 | if [[ ! -f tests/test.py ]]; then
5 | echo "Expected tests/test.py, are your running from repo root?"
6 | exit 1
7 | fi
8 | if [[ -d $TEST_DIR ]]; then
9 | echo "Exists! $TEST_DIR"
10 | exit 1
11 | fi
12 | mkdir -p $TEST_DIR
13 | pushd $TEST_DIR
14 | python3 -m venv .venv
15 | . ./.venv/bin/activate
16 | pip install torch numpy ninja torch pytest
17 | popd
18 | pip list
19 | python -m pytest tests/test.py
20 | deactivate
21 | rm -rf $TEST_DIR
22 |
23 |
--------------------------------------------------------------------------------
/logo/logo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fab-jul/torchac/63064ab4caeaa7d435d7dc6877160161066cb558/logo/logo.png
--------------------------------------------------------------------------------
/logo/logo_small.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fab-jul/torchac/63064ab4caeaa7d435d7dc6877160161066cb558/logo/logo_small.png
--------------------------------------------------------------------------------
/logo/logo_v1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fab-jul/torchac/63064ab4caeaa7d435d7dc6877160161066cb558/logo/logo_v1.png
--------------------------------------------------------------------------------
/logo/social.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fab-jul/torchac/63064ab4caeaa7d435d7dc6877160161066cb558/logo/social.png
--------------------------------------------------------------------------------
/pypi/test.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 | #
3 | # Test PyPi package
4 | #
5 |
6 |
7 | TESTS_PATH=$1
8 | PYTORCH_VERSION=$2
9 |
10 | if [[ -z $TESTS_PATH ]]; then
11 | echo "Usage: $0 TESTS_PATH [PYTORCH_VERSION]"
12 | exit 1
13 | fi
14 |
15 | if [[ -z $PYTORCH_VERSION ]]; then
16 | PYTORCH="pytorch" # Use latest
17 | else
18 | PYTORCH="pytorch==$PYTORCH_VERSION" # Use latest
19 | fi
20 |
21 | source ~/Documents/miniconda3/etc/profile.d/conda.sh
22 |
23 | ENV_NAME=torchac_test
24 | conda env list | grep $ENV_NAME
25 | if [[ $? == 1 ]]; then # Env does not exist yet.
26 | conda create -n $ENV_NAME pip python==3.8 -y
27 | fi
28 |
29 | set -e
30 | conda activate $ENV_NAME
31 | conda install $PYTORCH torchvision -c pytorch -y
32 | pip install pytest
33 | pip install --upgrade torchac --no-cache-dir
34 | python -c "import torchac"
35 | python -m pytest $TESTS_PATH -s
36 |
37 |
38 |
--------------------------------------------------------------------------------
/pypi/test_pytorch_versions.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | set -e
4 |
5 | source ~/Documents/miniconda3/etc/profile.d/conda.sh
6 |
7 | LOGFILE=tested.txt
8 | echo "Test Log" > $LOGFILE
9 |
10 | for VERSION in 1.5 1.6 1.7; do
11 | echo "Testing pytorch=$VERSION..."
12 | rm -rf ~/.cache/torch_extensions
13 | rm -rf data/MNIST # MNIST may lead to conflicts
14 | bash pypi/test.sh tests/test.py $VERSION
15 | echo "Tests pass for $VERSION" >> $LOGFILE
16 | conda activate torchac_test
17 | pip install matplotlib
18 | CUDA_VISIBLE_DEVICES="" python -u \
19 | examples/mnist_autoencoder/mnist_autoencoder_example.py \
20 | --max_training_itr 5
21 | conda deactivate
22 | conda env remove -n torchac_test -y
23 | echo "Example runs for $VERSION" >> $LOGFILE
24 | done
25 |
26 | echo ""
27 | cat $LOGFILE
28 | rm $LOGFILE
29 |
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | import setuptools
2 | import json
3 |
4 |
5 | def _get_long_description():
6 | with open('README.md', 'r') as f:
7 | long_description_lines = []
8 | skip = False
9 | for line in f:
10 | if '
5 | Output of the example script. First two columns show training 6 | set, second two columns show testing set.
7 |