tfc.entropy_models.UniversalBatchedEntropyModel | TensorFlow v2.16.1 (original) (raw)
tfc.entropy_models.UniversalBatchedEntropyModel
Stay organized with collections Save and categorize content based on your preferences.
Batched entropy model model which implements Universal Quantization.
View aliases
Main aliases
tfc.UniversalBatchedEntropyModel
tfc.entropy_models.UniversalBatchedEntropyModel(
prior,
coding_rank,
compression=False,
laplace_tail_mass=0.0,
expected_grads=False,
tail_mass=(2 ** -8),
range_coder_precision=12,
bottleneck_dtype=None,
num_noise_levels=15,
stateless=False,
decode_sanity_check=True
)
In contrast to the base class, which uses rounding for quantization, here "quantization" is performed additive uniform noise, which is implemented with Universal Quantization.
This is described in Sec. 3.2. in the paper
"Universally Quantized Neural Compression"
Eirikur Agustsson & Lucas Theis
https://arxiv.org/abs/2006.09952
| Args | |
|---|---|
| prior | A tfp.distributions.Distribution object. A density model fitting the marginal distribution of the bottleneck data with additive uniform noise, which is shared a priori between the sender and the receiver. For best results, the distribution should be flexible enough to have a unit-width uniform distribution as a special case, since this is the marginal distribution for bottleneck dimensions that are constant. The distribution parameters may not depend on data (they must be either variables or constants). |
| coding_rank | Integer. Number of innermost dimensions considered a coding unit. Each coding unit is compressed to its own bit string, and thebits() method sums over each coding unit. |
| compression | Boolean. If set to True, the range coding tables used bycompress() and decompress() will be built on instantiation. This assumes eager mode (throws an error if in graph mode or inside atf.function call). If set to False, these two methods will not be accessible. |
| laplace_tail_mass | Float. If positive, will augment the prior with a laplace mixture for training stability. |
| expected_grads | If True, will use analytical expected gradients during backpropagation w.r.t. additive uniform noise. |
| tail_mass | Float. Approximate probability mass which is encoded using an Elias gamma code embedded into the range coder. |
| range_coder_precision | Integer. Precision passed to the range coding op. |
| bottleneck_dtype | tf.dtypes.DType. Data type of bottleneck tensor. Defaults to tf.keras.mixed_precision.global_policy().compute_dtype. |
| num_noise_levels | Integer. The number of levels used to quantize the uniform noise. |
| stateless | Boolean. If True, creates range coding tables as Tensors rather than Variables. This makes the entropy model stateless and allows it to be constructed within a tf.function body. Ifcompression=False, then stateless=True is implied and the provided value is ignored. |
| decode_sanity_check | Boolean. If True, an raises an error if the binary strings passed into decompress are not completely decoded. |
| Attributes | |
|---|---|
| bottleneck_dtype | Data type of the bottleneck tensor. |
| cdf | The CDFs used by range coding. |
| cdf_offset | The CDF offsets used by range coding. |
| coding_rank | Number of innermost dimensions considered a coding unit. |
| compression | Whether this entropy model is prepared for compression. |
| expected_grads | Whether to use analytical expected gradients during backpropagation. |
| laplace_tail_mass | Whether to augment the prior with a NoisyLaplace mixture. |
| name | Returns the name of this module as passed or determined in the ctor. |
| name_scope | Returns a tf.name_scope instance for this class. |
| non_trainable_variables | Sequence of non-trainable variables owned by this module and its submodules. |
| prior | Prior distribution, used for deriving range coding tables. |
| prior_shape | Batch shape of prior (dimensions which are not assumed i.i.d.). |
| prior_shape_tensor | Batch shape of prior as a Tensor. |
| range_coder_precision | Precision used in range coding op. |
| stateless | Whether range coding tables are created as Tensors or Variables. |
| submodules | Sequence of all sub-modules.Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on). a = tf.Module() b = tf.Module() c = tf.Module() a.b = b b.c = c list(a.submodules) == [b, c] True list(b.submodules) == [c] True list(c.submodules) == [] True |
| tail_mass | Approximate probability mass which is range encoded with overflow. |
| trainable_variables | Sequence of trainable variables owned by this module and its submodules. |
| variables | Sequence of variables owned by this module and its submodules. |
Methods
compress
compress(
bottleneck
)
Compresses a floating-point tensor.
Compresses the tensor to bit strings. bottleneck is first quantized as in quantize(), and then compressed using the probability tables inself.cdf derived from self.prior. The quantized tensor can later be recovered by calling decompress().
The innermost self.coding_rank dimensions are treated as one coding unit, i.e. are compressed into one string each. Any additional dimensions to the left are treated as batch dimensions.
| Args | |
|---|---|
| bottleneck | tf.Tensor containing the data to be compressed. Must have at least self.coding_rank dimensions, and the innermost dimensions must be broadcastable to self.prior_shape. |
| Returns |
|---|
| A tf.Tensor having the same shape as bottleneck without theself.coding_rank innermost dimensions, containing a string for each coding unit. |
decompress
decompress(
strings, broadcast_shape
)
Decompresses a tensor.
Reconstructs the quantized tensor from bit strings produced by compress(). It is necessary to provide a part of the output shape in broadcast_shape.
| Args | |
|---|---|
| strings | tf.Tensor containing the compressed bit strings. |
| broadcast_shape | Iterable of ints. The part of the output tensor shape between the shape of strings on the left and self.prior_shape on the right. This must match the shape of the input to compress(). |
| Returns |
|---|
| A tf.Tensor of shape strings.shape + broadcast_shape + self.prior_shape. |
get_config
get_config()
Returns the configuration of the entropy model.
| Returns |
|---|
| A JSON-serializable Python dict. |
| Raises | |
|---|---|
| RuntimeError | on attempting to call this method on an entropy model with compression=False or with stateless=True. |
get_weights
get_weights()
set_weights
set_weights(
weights
)
with_name_scope
@classmethod
with_name_scope( method )
Decorator to automatically enter the module name scope.
class MyModule(tf.Module):
@tf.Module.with_name_scope
def __call__(self, x):
if not hasattr(self, 'w'):
self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
return tf.matmul(x, self.w)
Using the above module would produce tf.Variables and tf.Tensors whose names included the module name:
mod = MyModule()
mod(tf.ones([1, 2]))
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
mod.w
<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
numpy=..., dtype=float32)>
| Args | |
|---|---|
| method | The method to wrap. |
| Returns |
|---|
| The original method wrapped such that it enters the module's name scope. |
__call__
__call__(
bottleneck, training=True
)
Perturbs a tensor with additive uniform noise and estimates bitcost.
| Args | |
|---|---|
| bottleneck | tf.Tensor containing a non-perturbed bottleneck. Must have at least self.coding_rank dimensions. |
| training | Boolean. If False, computes the bitcost using discretized uniform noise. If True, estimates the differential entropy with uniform noise. |
| Returns |
|---|
| A tuple (bottleneck_perturbed, bits) where bottleneck_perturbed is bottleneck perturbed with nosie and bits is the bitcost of transmitting such a sample having the same shape as bottleneck without the self.coding_rank innermost dimensions. |