|
| 1 | +# Lint as: python3 |
| 2 | +# Copyright 2020 Google LLC. All Rights Reserved. |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | +# ============================================================================== |
| 16 | +"""Base class for continuous entropy models.""" |
| 17 | + |
| 18 | +import abc |
| 19 | + |
| 20 | +from absl import logging |
| 21 | +import tensorflow.compat.v2 as tf |
| 22 | + |
| 23 | +from tensorflow_compression.python.distributions import helpers |
| 24 | +from tensorflow_compression.python.ops import range_coding_ops |
| 25 | + |
| 26 | + |
| 27 | +__all__ = ["ContinuousEntropyModelBase"] |
| 28 | + |
| 29 | + |
| 30 | +class ContinuousEntropyModelBase(tf.Module, metaclass=abc.ABCMeta): |
| 31 | + """Base class for continuous entropy models. |
| 32 | +
|
| 33 | + The basic functionality of this class is to pre-compute integer probability |
| 34 | + tables based on the provided `tfp.distributions.Distribution` object, which |
| 35 | + can then be used reliably across different platforms by the range coder and |
| 36 | + decoder. |
| 37 | + """ |
| 38 | + |
| 39 | + @abc.abstractmethod |
| 40 | + def __init__(self, distribution, coding_rank, |
| 41 | + likelihood_bound=1e-9, tail_mass=2**-8, |
| 42 | + range_coder_precision=12): |
| 43 | + """Initializer. |
| 44 | +
|
| 45 | + Arguments: |
| 46 | + distribution: A `tfp.distributions.Distribution` object modeling the |
| 47 | + distribution of the input data including additive uniform noise. For |
| 48 | + best results, the distribution should be flexible enough to have a |
| 49 | + unit-width uniform distribution as a special case. |
| 50 | + coding_rank: Integer. Number of innermost dimensions considered a coding |
| 51 | + unit. Each coding unit is compressed to its own bit string, and the |
| 52 | + `bits()` method sums over each coding unit. |
| 53 | + likelihood_bound: Float. Lower bound for likelihood values, to prevent |
| 54 | + training instabilities. |
| 55 | + tail_mass: Float. Approximate probability mass which is range encoded with |
| 56 | + less precision, by using a Golomb-like code. |
| 57 | + range_coder_precision: Integer. Precision passed to the range coding op. |
| 58 | + """ |
| 59 | + if not distribution.is_scalar_event(): |
| 60 | + raise ValueError( |
| 61 | + "`distribution` must be a (batch of) scalar distribution(s).") |
| 62 | + super().__init__() |
| 63 | + self._distribution = distribution |
| 64 | + self._coding_rank = int(coding_rank) |
| 65 | + self._likelihood_bound = float(likelihood_bound) |
| 66 | + self._tail_mass = float(tail_mass) |
| 67 | + self._range_coder_precision = int(range_coder_precision) |
| 68 | + self.update_tables() |
| 69 | + |
| 70 | + @property |
| 71 | + def distribution(self): |
| 72 | + """Distribution modeling data + i.i.d. uniform noise.""" |
| 73 | + return self._distribution |
| 74 | + |
| 75 | + @property |
| 76 | + def coding_rank(self): |
| 77 | + """Number of innermost dimensions considered a coding unit.""" |
| 78 | + return self._coding_rank |
| 79 | + |
| 80 | + @property |
| 81 | + def likelihood_bound(self): |
| 82 | + """Lower bound for likelihood values.""" |
| 83 | + return self._likelihood_bound |
| 84 | + |
| 85 | + @property |
| 86 | + def tail_mass(self): |
| 87 | + """Approximate probability mass which is range encoded with overflow.""" |
| 88 | + return self._tail_mass |
| 89 | + |
| 90 | + @property |
| 91 | + def range_coder_precision(self): |
| 92 | + """Precision passed to range coding op.""" |
| 93 | + return self._range_coder_precision |
| 94 | + |
| 95 | + @property |
| 96 | + def dtype(self): |
| 97 | + """Data type of this distribution.""" |
| 98 | + return self.distribution.dtype |
| 99 | + |
| 100 | + def quantization_offset(self): |
| 101 | + """Distribution-dependent quantization offset.""" |
| 102 | + return helpers.quantization_offset(self.distribution) |
| 103 | + |
| 104 | + def lower_tail(self): |
| 105 | + """Approximate lower tail quantile for range coding.""" |
| 106 | + return helpers.lower_tail(self.distribution, self.tail_mass) |
| 107 | + |
| 108 | + def upper_tail(self): |
| 109 | + """Approximate upper tail quantile for range coding.""" |
| 110 | + return helpers.upper_tail(self.distribution, self.tail_mass) |
| 111 | + |
| 112 | + @tf.custom_gradient |
| 113 | + def _quantize(self, inputs, offset): |
| 114 | + return tf.round(inputs - offset) + offset, lambda x: (x, None) |
| 115 | + |
| 116 | + def update_tables(self): |
| 117 | + """Updates integer-valued probability tables used by the range coder. |
| 118 | +
|
| 119 | + These tables must not be re-generated independently on the sending and |
| 120 | + receiving side, since small numerical discrepancies between both sides can |
| 121 | + occur in this process. If the tables differ slightly, this in turn would |
| 122 | + very likely cause catastrophic error propagation during range decoding. For |
| 123 | + a more in-depth discussion of this, see: |
| 124 | +
|
| 125 | + > "Integer Networks for Data Compression with Latent-Variable Models"<br /> |
| 126 | + > J. Ballé, N. Johnston, D. Minnen<br /> |
| 127 | + > https://openreview.net/forum?id=S1zz2i0cY7 |
| 128 | +
|
| 129 | + The tables are stored in `tf.Tensor`s as attributes of this object. The |
| 130 | + recommended way is to train the model, then call this method, and then |
| 131 | + distribute the model to a sender and a receiver. |
| 132 | + """ |
| 133 | + offset = self.quantization_offset() |
| 134 | + lower_tail = self.lower_tail() |
| 135 | + upper_tail = self.upper_tail() |
| 136 | + |
| 137 | + # Largest distance observed between lower tail and median, and between |
| 138 | + # median and upper tail. |
| 139 | + minima = offset - lower_tail |
| 140 | + minima = tf.cast(tf.math.ceil(minima), tf.int32) |
| 141 | + minima = tf.math.maximum(minima, 0) |
| 142 | + maxima = upper_tail - offset |
| 143 | + maxima = tf.cast(tf.math.ceil(maxima), tf.int32) |
| 144 | + maxima = tf.math.maximum(maxima, 0) |
| 145 | + |
| 146 | + # PMF starting positions and lengths. |
| 147 | + pmf_start = offset - tf.cast(minima, self.dtype) |
| 148 | + pmf_length = maxima + minima + 1 |
| 149 | + |
| 150 | + # Sample the densities in the computed ranges, possibly computing more |
| 151 | + # samples than necessary at the upper end. |
| 152 | + max_length = tf.math.reduce_max(pmf_length) |
| 153 | + if max_length > 2048: |
| 154 | + logging.warning( |
| 155 | + "Very wide PMF with %d elements may lead to out of memory issues. " |
| 156 | + "Consider encoding distributions with smaller dispersion or " |
| 157 | + "increasing `tail_mass` parameter.", int(max_length)) |
| 158 | + samples = tf.range(tf.cast(max_length, self.dtype), dtype=self.dtype) |
| 159 | + samples = tf.reshape( |
| 160 | + samples, [-1] + self.distribution.batch_shape.rank * [1]) |
| 161 | + samples += pmf_start |
| 162 | + pmf = self.distribution.prob(samples) |
| 163 | + |
| 164 | + # Collapse batch dimensions of distribution. |
| 165 | + pmf = tf.reshape(pmf, [max_length, -1]) |
| 166 | + pmf = tf.transpose(pmf) |
| 167 | + |
| 168 | + dist_shape = self.distribution.batch_shape_tensor() |
| 169 | + pmf_length = tf.broadcast_to(pmf_length, dist_shape) |
| 170 | + pmf_length = tf.reshape(pmf_length, [-1]) |
| 171 | + cdf_length = pmf_length + 2 |
| 172 | + cdf_offset = tf.broadcast_to(-minima, dist_shape) |
| 173 | + cdf_offset = tf.reshape(cdf_offset, [-1]) |
| 174 | + |
| 175 | + # Prevent tensors from bouncing back and forth between host and GPU. |
| 176 | + with tf.device("/cpu:0"): |
| 177 | + def loop_body(args): |
| 178 | + prob, length = args |
| 179 | + prob = prob[:length] |
| 180 | + prob = tf.concat([prob, 1 - tf.reduce_sum(prob, keepdims=True)], axis=0) |
| 181 | + cdf = range_coding_ops.pmf_to_quantized_cdf( |
| 182 | + prob, precision=self.range_coder_precision) |
| 183 | + return tf.pad( |
| 184 | + cdf, [[0, max_length - length]], mode="CONSTANT", constant_values=0) |
| 185 | + |
| 186 | + # TODO(jonycgn,ssjhv): Consider switching to Python control flow. |
| 187 | + cdf = tf.map_fn( |
| 188 | + loop_body, (pmf, pmf_length), dtype=tf.int32, name="pmf_to_cdf") |
| 189 | + |
| 190 | + self._cdf, self._cdf_offset, self._cdf_length = cdf, cdf_offset, cdf_length |
0 commit comments