def test_tails_and_offset_are_in_order(self): dist = self.dist_cls(loc=10.3, scale=1.5) offset = helpers.quantization_offset(dist) lower_tail = helpers.lower_tail(dist, 2**-8) upper_tail = helpers.upper_tail(dist, 2**-8) self.assertGreater(upper_tail, lower_tail) self.assertAllClose(offset, 0.3)
def test_tails_and_offset_are_in_order(self):
dist = self.dist_cls(loc=10, scale=[1.5, 2], weight=[.5, .5])
offset = helpers.quantization_offset(dist)
lower_tail = helpers.lower_tail(dist, 2**-8)
upper_tail = helpers.upper_tail(dist, 2**-8)
self.assertGreater(upper_tail, offset)
self.assertGreater(offset, lower_tail)
def test_tails_and_offset_are_in_order(self): dist = self.dist_cls(loc=[5.4, 8.6], scale=[1.4, 2], weight=[.6, .4]) offset = helpers.quantization_offset(dist) lower_tail = helpers.lower_tail(dist, 2**-8) upper_tail = helpers.upper_tail(dist, 2**-8) self.assertGreater(upper_tail, lower_tail) self.assertAllClose(offset, 0.4) # Decimal part of the peakiest mode (5.4).
def test_tails_and_offset_are_in_order(self):
df = deep_factorized.NoisyDeepFactorized()
offset = helpers.quantization_offset(df)
lower_tail = helpers.lower_tail(df, 2**-8)
upper_tail = helpers.upper_tail(df, 2**-8)
self.assertGreater(upper_tail, offset)
self.assertGreater(offset, lower_tail)
def _build_tables(self, prior):
"""Computes integer-valued probability tables used by the range coder.
These tables must not be re-generated independently on the sending and
receiving side, since small numerical discrepancies between both sides can
occur in this process. If the tables differ slightly, this in turn would
very likely cause catastrophic error propagation during range decoding. For
a more in-depth discussion of this, see:
> "Integer Networks for Data Compression with Latent-Variable Models"<br />
> J. Ballé, N. Johnston, D. Minnen<br />
> https://openreview.net/forum?id=S1zz2i0cY7
The tables are stored in `tf.Variable`s as attributes of this object. The
recommended way is to train the model, instantiate an entropy model with
`compression=True`, and then distribute the model to a sender and a
receiver.
Args:
prior: The `tfp.distributions.Distribution` object (see initializer).
"""
# TODO(jonycgn, relational): Consider not using offset when soft quantization
# is used.
offset = self._offset_from_prior(prior)
lower_tail = helpers.lower_tail(prior, self.tail_mass)
upper_tail = helpers.upper_tail(prior, self.tail_mass)
# Integers such that:
# minima + offset < lower_tail
# maxima + offset > upper_tail
minima = tf.cast(tf.math.floor(lower_tail - offset), tf.int32)
maxima = tf.cast(tf.math.ceil(upper_tail - offset), tf.int32)
# PMF starting positions and lengths.
pmf_start = tf.cast(minima, self.dtype) + offset
pmf_length = maxima - minima + 1
# Sample the densities in the computed ranges, possibly computing more
# samples than necessary at the upper end.
max_length = tf.math.reduce_max(pmf_length)
if tf.executing_eagerly() and max_length > 2048:
logging.warning(
"Very wide PMF with %d elements may lead to out of memory issues. "
"Consider priors with smaller dispersion or increasing `tail_mass` "
"parameter.", int(max_length))
samples = tf.range(tf.cast(max_length, self.dtype), dtype=self.dtype)
samples = tf.reshape(samples, [-1] + len(self.context_shape) * [1])
samples += pmf_start
pmf = prior.prob(samples)
# Collapse batch dimensions of distribution.
pmf = tf.reshape(pmf, [max_length, -1])
pmf = tf.transpose(pmf)
pmf_length = tf.broadcast_to(pmf_length, self.context_shape_tensor)
pmf_length = tf.reshape(pmf_length, [-1])
cdf_length = pmf_length + 2
cdf_offset = tf.broadcast_to(minima, self.context_shape_tensor)
cdf_offset = tf.reshape(cdf_offset, [-1])
# Prevent tensors from bouncing back and forth between host and GPU.
with tf.device("/cpu:0"):
def loop_body(args):
prob, length = args
prob = prob[:length]
overflow = tf.math.maximum(1 - tf.reduce_sum(prob, keepdims=True), 0.)
prob = tf.concat([prob, overflow], axis=0)
cdf = range_coding_ops.pmf_to_quantized_cdf(
prob, precision=self.range_coder_precision)
return tf.pad(
cdf, [[0, max_length - length]], mode="CONSTANT", constant_values=0)
# TODO(jonycgn,ssjhv): Consider switching to Python control flow.
cdf = tf.map_fn(
loop_body, (pmf, pmf_length), dtype=tf.int32, name="pmf_to_cdf")
if self.no_variables:
self._cdf = cdf
self._cdf_offset = cdf_offset
self._cdf_length = cdf_length
else:
self._cdf = tf.Variable(cdf, trainable=False, name="cdf")
self._cdf_offset = tf.Variable(
cdf_offset, trainable=False, name="cdf_offset")
self._cdf_length = tf.Variable(
cdf_length, trainable=False, name="cdf_length")
def test_tails_are_in_order(self):
df = deep_factorized.NoisyDeepFactorized()
lower_tail = helpers.lower_tail(df, 2**-8)
upper_tail = helpers.upper_tail(df, 2**-8)
self.assertGreater(upper_tail, lower_tail)
def test_deep_factorized_tails_are_in_order(self):
dist = deep_factorized.DeepFactorized(batch_shape=[10])
self.assertAllGreater(
helpers.upper_tail(dist, 2**-8) - helpers.lower_tail(dist, 2**-8),
0)
def test_normal_tails_are_in_order(self):
dist = tfp.distributions.Normal(loc=3., scale=5.)
self.assertGreater(helpers.upper_tail(dist, 2**-8),
helpers.lower_tail(dist, 2**-8))
def test_logistic_tails_are_in_order(self):
dist = tfp.distributions.Logistic(loc=-3., scale=1.)
self.assertGreater(helpers.upper_tail(dist, 2**-8),
helpers.lower_tail(dist, 2**-8))
def test_laplace_tails_are_in_order(self):
dist = tfp.distributions.Laplace(loc=-2., scale=5.)
self.assertGreater(helpers.upper_tail(dist, 2**-8),
helpers.lower_tail(dist, 2**-8))
def test_cauchy_tails_are_in_order(self):
dist = tfp.distributions.Cauchy(loc=1.5, scale=3.)
self.assertGreater(helpers.upper_tail(dist, 2**-8),
helpers.lower_tail(dist, 2**-8))
def _build_tables(self, prior, precision, offset=None):
"""Computes integer-valued probability tables used by the range coder.
These tables must not be re-generated independently on the sending and
receiving side, since small numerical discrepancies between both sides can
occur in this process. If the tables differ slightly, this in turn would
very likely cause catastrophic error propagation during range decoding. For
a more in-depth discussion of this, see:
> "Integer Networks for Data Compression with Latent-Variable Models"<br />
> J. Ballé, N. Johnston, D. Minnen<br />
> https://openreview.net/forum?id=S1zz2i0cY7
Args:
prior: The `tfp.distributions.Distribution` object (see initializer).
precision: Integer. Precision for range coder.
offset: None or float tensor between -.5 and +.5. Sub-integer quantization
offsets to use for sampling prior probabilities. Defaults to 0.
Returns:
CDF table, CDF offsets, CDF lengths.
"""
precision = int(precision)
offset = tf.cast(0 if offset is None else offset, prior.dtype)
# Subclasses should have already caught this, but better be safe.
assert not prior.event_shape.rank
lower_tail = helpers.lower_tail(prior, self.tail_mass)
upper_tail = helpers.upper_tail(prior, self.tail_mass)
# Integers such that:
# minima + offset < lower_tail
# maxima + offset > upper_tail
minima = tf.cast(tf.math.floor(lower_tail - offset), tf.int32)
maxima = tf.cast(tf.math.ceil(upper_tail - offset), tf.int32)
# PMF starting positions and lengths.
pmf_start = tf.cast(minima, prior.dtype) + offset
pmf_length = maxima - minima + 1
# Sample the densities in the computed ranges, possibly computing more
# samples than necessary at the upper end.
max_length = tf.math.reduce_max(pmf_length)
if tf.executing_eagerly() and max_length > 2048:
logging.warning(
"Very wide PMF with %d elements may lead to out of memory issues. "
"Consider priors with smaller variance, or increasing `tail_mass` "
"parameter.", int(max_length))
samples = tf.range(tf.cast(max_length, prior.dtype), dtype=prior.dtype)
samples = tf.reshape(samples, [-1] + pmf_length.shape.rank * [1])
samples += pmf_start
pmf = prior.prob(samples)
pmf_shape = tf.shape(pmf)[1:]
num_pmfs = tf.reduce_prod(pmf_shape)
# Collapse batch dimensions of distribution.
pmf = tf.reshape(pmf, [max_length, num_pmfs])
pmf = tf.transpose(pmf)
pmf_length = tf.broadcast_to(pmf_length, pmf_shape)
pmf_length = tf.reshape(pmf_length, [num_pmfs])
cdf_offset = tf.broadcast_to(minima, pmf_shape)
cdf_offset = tf.reshape(cdf_offset, [num_pmfs])
precision_tensor = tf.constant([-precision], dtype=tf.int32)
# Prevent tensors from bouncing back and forth between host and GPU.
with tf.device("/cpu:0"):
def loop_body(i, cdf):
p = pmf[i, :pmf_length[i]]
overflow = tf.math.maximum(1. - tf.reduce_sum(p, keepdims=True), 0.)
p = tf.cast(tf.concat([p, overflow], 0), tf.float32)
c = gen_ops.pmf_to_quantized_cdf(p, precision=precision)
return i + 1, tf.concat([cdf, precision_tensor, c], 0)
i_0 = tf.constant(0, tf.int32)
cdf_0 = tf.constant([], tf.int32)
_, cdf = tf.while_loop(
lambda i, _: i < num_pmfs, loop_body, (i_0, cdf_0),
shape_invariants=(i_0.shape, tf.TensorShape([None])),
name="pmf_to_cdf")
return cdf, cdf_offset
def _upper_tail(self, tail_mass):
return tf.math.ceil(helpers.upper_tail(self.base, tail_mass))
def _upper_tail(self, tail_mass):
# Same logic as for _quantile.
if not self.invertible:
raise NotImplementedError()
return self.transform(helpers.upper_tail(self.base, tail_mass))
def upper_tail(self):
"""Approximate upper tail quantile for range coding."""
return helpers.upper_tail(self.prior, self.tail_mass)
def _upper_tail(self, tail_mass):
return helpers.upper_tail(self.base, tail_mass)