Quantizer

quantize(
x: Tensor,
quantizer_noise_type: Literal['kumaraswamy', 'gaussian', 'none'] = 'kumaraswamy',
quantizer_type: Literal['softround_alone', 'softround', 'hardround', 'ste', 'none'] = 'softround',
soft_round_temperature: float | None = 0.3,
noise_parameter: float | None = 1.0,
) → Tensor[source]

Quantize an input \(x\) to an output \(y\) simulating the quantization. There is different mode possibles, described by quantizer_type:

  • none: \(y = x + n\) with \(n\) a random noise (more details below)

  • softround_alone: \(y = \mathrm{softround}(x, t)\) with \(t\) the soft_round_temperature.

  • softround: \(y = \mathrm{softround}(\mathrm{softround}(x, t) + n, t)\) with \(t\) the soft_round_temperature and \(n\) a random noise (more details below)

  • hardround: \(y = \mathrm{round}(x)\)

  • ste: \(y = \mathrm{round}(x)\) (backward done through softround)

The noise is parameterized by quantizer_noise_type and noise_parameter. This last parameter has a different role for the different noise type:

  • gaussian: noise_parameter is the standard deviation of the gaussian distribution

  • kumaraswamy: noise_parameter corresponds to the \(a\) parameter of the kumaraswamy distribution. 1 means uniform distribution and increasing it leads to more more and more probability of being into the center.

Softround is parameterized by soft_round_temperature denoted as \(t\). Setting \(t = 0\) corresponds to the actual quantization i.e. round(x). As \(t\) grows bigger, the function approaches identity i.e. \(\lim_{t \rightarrow \infty} \mathrm{softround}(x, t) = x\). In practice \(t \geq 1\) is already quite close to identity.,

Note

Why do we apply twice the softround when quantizer_type is softround? It follows the operations described in C3: High-performance and low-complexity neural compression from a single image or video, Kim et al. i.e.

  1. Use a soft round function instead of the non-differentiable round function

  2. Add a random noise to prevent the network from learning the inverse softround function

  3. Re-apply the soft round function as advocated in Universally Quantized Neural Compression, Agustsson & Theis

Parameters:

  • x (Tensor) – Tensor to be quantized.

  • quantizer_noise_type (Literal['kumaraswamy', 'gaussian', 'none']) – noise type. Defaults to "kumaraswamy".

  • quantizer_type (Literal['softround_alone', 'softround', 'hardround', 'ste', 'none']) – quantizer type. Defaults to "softround".

  • soft_round_temperature (float | None) – Soft round temperature. This is used for softround modes as well as the ste mode to simulate the derivative in the backward. Defaults to 0.3.

  • noise_parameter (float | None) – noise distribution parameter. Defaults to 1.0.

Returns:

Quantized tensor

Return type:

Tensor

softround(x: Tensor, t: float) → Tensor[source]

Perform the softround function as introduced in section 4.1 of the paper Universally Quantized Neural Compression, Agustsson & Theis, defined as follows:

\[\mathrm{softround}(x, t) = \lfloor x \rfloor + \frac{\mathrm{tanh}(\frac{\Delta}{t})}{2\ \mathrm{\mathrm{tanh}(\frac{1}{2t})}} + \frac{1}{2}, \text{ with } \Delta = x - \lfloor x \rfloor - \frac{1}{2}.\]
Parameters:

  • x (Tensor) – Input tensor to be quantized.

  • t (float) – Soft round temperature \(t\). Setting \(t = 0\) corresponds to the actual quantization i.e. round(x). As \(t\) grows bigger, the function approaches identity i.e. \(\lim_{t \rightarrow \infty} \mathrm{softround}(x, t) = x\). In practice \(t \geq 1\) is already quite close to identity.

Returns:

Soft-rounded tensor

Return type:

Tensor

generate_kumaraswamy_noise(uniform_noise: Tensor, kumaraswamy_param: float) → Tensor[source]

Reparameterize a random variable uniform_noise following a uniform distribution \(\mathcal{U}(0, 1)\) to a random variable following a kumaraswamy distribution as proposed in the paper C3: High-performance and low-complexity neural compression from a single image or video, Kim et al.

The kumaraswamy distribution is defined on the interval \((0, 1)\) with the following PDF:

\[f(x;a,b) = 1 - (1 - x^a)^b\]

Here, it is only parameterized through a single parameter kumaraswamy_param corresponding to \(a\) in the above equation. The second parameter \(b\) is set to as a function of \(a\) so that the mode of the distribution is always \(\frac{1}{2}\). Setting \(a=1\) gives the uniform distribution \(\mathcal{U}(0, 1)\). Increasing the value of \(a\) gives more “pointy” distribution.

The resulting kumaraswamy noise is shifted so that it lies in \((-\frac{1}{2}, \frac{1}{2})\).

Parameters:

  • uniform_noise (Tensor) – A uniform noise in \([0, 1]\) with any size.

  • kumaraswamy_param (float) – Parameter \(a\) of a Kumaraswamy distribution. Set it to 1 for a uniform noise.

Returns:

A kumaraswamy noise with identical dim to uniform_noise in \([-\frac{1}{2}, \frac{1}{2}]\).

Return type:

Tensor