espnet2.gan_codec.shared.quantizer.modules.core_vq.EuclideanCodebook

source

class espnet2.gan_codec.shared.quantizer.modules.core_vq.EuclideanCodebook(dim: int, codebook_size: int, kmeans_init: int = False, kmeans_iters: int = 10, decay: float = 0.99, epsilon: float = 1e-05, threshold_ema_dead_code: int = 2)

Bases: Module

Codebook with Euclidean distance. :param dim: Dimension. :type dim: int :param codebook_size: Codebook size. :type codebook_size: int :param kmeans_init: Whether to use k-means to initialize the codebooks.

If set to true, run the k-means algorithm on the first training batch and use the learned centroids as initialization.

• Parameters:
  • kmeans_iters (int) – Number of iterations used for k-means algorithm at initialization.
  • decay (float) – Decay for exponential moving average over the codebooks.
  • epsilon (float) – Epsilon value for numerical stability.
  • threshold_ema_dead_code (int) – Threshold for dead code expiration. Replace any codes that have an exponential moving average cluster size less than the specified threshold with randomly selected vector from the current batch.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

decode(embed_ind)

dequantize(embed_ind)

encode(x)

expire_codes_(batch_samples)

forward(x)

Codebook Forward with EMA.

• Parameters:x (Tensor) – Vector for quantization (B, T, D)
• Returns: Quantized output (B, T, D) Tensor: Codebook Index (B, T)
• Return type: Tensor

init_embed_(data)

postprocess_emb(embed_ind, shape)

preprocess(x)

quantize(x)

replace_(samples, mask)