speechbrain.nnet.quantisers 模块

Gumbel Softmax 实现，支持多个组。

Authors

• 鲁道夫·A·布劳恩 2022

摘要

类：

GumbelVectorQuantizer	使用Gumbel softmax进行向量量化。
RandomProjectionQuantizer	使用投影和随机初始化的码本进行向量量化，这对于像BEST-RQ这样的模型非常有用。

参考

class speechbrain.nnet.quantisers.GumbelVectorQuantizer(input_dim, num_vars, temp_tuple, groups, vq_dim)

基础：Module

使用Gumbel Softmax进行向量量化。复制自fairseq实现。 :param input_dim: 输入维度（通道数）。 :type input_dim: int :param num_vars: 每组量化向量的数量。 :type num_vars: int :param temp_tuple: 训练时的温度。这应该是一个包含3个元素的元组：（开始，停止，衰减因子）。 :type temp_tuple: float :param groups: 向量量化的组数。 :type groups: int :param vq_dim: 结果量化向量的维度。 :type vq_dim: int

Example

>>> quantiser = GumbelVectorQuantizer(128, 100, (2.0, 0.25, 0.999995,), 2, 50 )
>>> inputs = torch.rand(10, 12, 128)
>>> output = quantiser(inputs)
>>> output["x"].shape
torch.Size([10, 12, 50])

update_temp(steps)

根据当前步骤更新温度

forward(x)

将潜在向量前向传播以获得量化输出

class speechbrain.nnet.quantisers.RandomProjectionQuantizer(input_dim, cb_dim, cb_vocab)

基础：Module

使用投影和随机初始化的码本进行向量量化 这对于像BEST-RQ这样的模型非常有用。

输出是输入每个时间步的码本中最接近代码的索引。

参考: https://arxiv.org/pdf/2202.01855

Parameters:

• input_dim (int) – 输入维度（通道数）。

• cb_dim (int) – 码本中每个代码的大小。

• cb_vocab (int) – 码本中的代码数量

Example

>>> quantiser = RandomProjectionQuantizer(16, 16, 32)
>>> inputs = torch.rand(10, 12, 16)
>>> output = quantiser(inputs)
>>> output.shape
torch.Size([10, 12])

forward(x)

将潜在向量前向传播以获得量化输出