espnet2.asr.encoder.beats_encoder.quant_noise

source

espnet2.asr.encoder.beats_encoder.quant_noise(module, p, block_size)

Wraps modules and applies quantization noise to the weights for subsequent quantization with Iterative Product Quantization, as described in “Training with Quantization Noise for Extreme Model Compression”.

This function modifies the behavior of a given module by introducing quantization noise during training. It randomly drops blocks of weights in the module, which can help improve the robustness of the model to quantization effects.

• Parameters:
  • module (nn.Module) – The PyTorch module (e.g., Linear, Embedding, or Conv2d) to which quantization noise will be applied.
  • p (float) – The probability of dropping a block of weights. Should be a value between 0 and 1.
  • block_size (int) – The size of the blocks for subsequent quantization with iterative product quantization.
• Returns: The modified module with quantization noise applied.
• Return type: nn.Module
• Raises:AssertionError – If the module type is not supported or if the weights of the module do not have the correct dimensions with respect to the specified block size.

NOTE

• Module weights must have the right sizes relative to the block size.
• Only Linear, Embedding, and Conv2d modules are supported.
• For more details on how to quantize by blocks with convolutional weights, see “And the Bit Goes Down: Revisiting the Quantization of Neural Networks”.
• This implementation represents a simple form of noise, which consists of randomly dropping blocks.

Examples

>>> linear_layer = nn.Linear(10, 5)
>>> noisy_layer = quant_noise(linear_layer, p=0.1, block_size=2)
>>> print(noisy_layer)