speechbrain.lobes.models.discrete.dac 模块

此 lobe 支持集成预训练的离散 DAC 模型。参考： http://arxiv.org/abs/2306.06546 参考： https://descript.notion.site/Descript-Audio-Codec-11389fce0ce2419891d6591a68f814d5 参考： https://github.com/descriptinc/descript-audio-codec

作者

Shubham Gupta 2023

摘要

类

`DAC`	用于音频数据编码和解码的离散自编码器编解码器 (DAC)。
`Decoder`	用于 DAC 解码器部分的 PyTorch 模块。
`DecoderBlock`	表示解码器架构中一个块的 PyTorch 模块。
`Encoder`	用于 DAC 编码器部分的 PyTorch 模块。
`EncoderBlock`	用于卷积神经网络的编码器块模块。
`ResidualUnit`	用于卷积神经网络的残差单元模块。
`ResidualVectorQuantize`	在 SoundStream 中引入：一个端到端神经音频编解码器 https://arxiv.org/abs/2107.03312
`Snake1d`	实现在 1D 中 Snake 激活函数的 PyTorch 模块。
`VectorQuantize`	向量量化实现

函数

`WNConv1d`	对 1D 卷积层应用权重归一化。
`WNConvTranspose1d`	对 1D 转置卷积层应用权重归一化。
`download`	根据模型类型、比特率和标签下载指定的模型文件，并将其保存到本地路径。
`init_weights`	初始化 1D 卷积层的权重。

参考

speechbrain.lobes.models.discrete.dac.WNConv1d(*args, **kwargs)[source]

对 1D 卷积层应用权重归一化。

参数:

*args (tuple) – nn.Conv1d 的可变长度参数列表。
**kwargs (dict) – nn.Conv1d 的任意关键字参数。

返回:

权重归一化的 nn.Conv1d 层。

返回类型:

torch.nn.Module

speechbrain.lobes.models.discrete.dac.WNConvTranspose1d(*args, **kwargs)[source]

对 1D 转置卷积层应用权重归一化。

参数:

*args (tuple) – nn.ConvTranspose1d 的可变长度参数列表。
**kwargs (dict) – nn.ConvTranspose1d 的任意关键字参数。

返回:

权重归一化的 nn.ConvTranspose1d 层。

返回类型:

torch.nn.Module

speechbrain.lobes.models.discrete.dac.init_weights(m)[source]: 初始化 1D 卷积层的权重。

speechbrain.lobes.models.discrete.dac.download(model_type: str = '44khz', model_bitrate: str = '8kbps', tag: str = 'latest', local_path: Path | None = None)[source]

根据模型类型、比特率和标签下载指定的模型文件，并将其保存到本地路径。

参数:

model_type (str, 可选) – 要下载的模型类型。可以是 ‘44khz’、‘24khz’ 或 ‘16khz’。默认为 ‘44khz’。
model_bitrate (str, 可选) – 模型的比特率。可以是 ‘8kbps’ 或 ‘16kbps’。默认为 ‘8kbps’。
tag (str, 可选) – 模型的特定版本标签。默认为 ‘latest’。
local_path (Path, 可选) – 模型将保存到的本地文件路径。如果未提供，将使用默认路径。

返回:

模型保存到的本地路径。

返回类型:

Path

引发:

ValueError – 如果模型类型或比特率不受支持，或者模型无法找到或下载。

class speechbrain.lobes.models.discrete.dac.VectorQuantize(input_dim: int, codebook_size: int, codebook_dim: int)[source]

基类: Module

向量量化实现

VQ 的实现，类似于 Karpathy 的 repo： https://github.com/karpathy/deep-vector-quantization 此外，还使用了 Improved VQGAN (https://arxiv.org/pdf/2110.04627.pdf) 中的以下技巧：

分解码：在低维空间中执行最近邻查找
以改善码本使用率

l2 归一化码：将欧几里得距离转换为余弦相似度，这
提高了训练稳定性

参数:

input_dim (int) – 输入的维度
codebook_size (int) – 码本大小
codebook_dim (int) – 码本的维度

forward(z: Tensor)[source]

使用固定码本对输入张量进行量化，并返回相应的码本向量

参数:

z (torch.Tensor[B x D x T])

返回:

torch.Tensor[B x D x T] – 输入的量化连续表示
torch.Tensor[1] – 用于训练编码器以预测更接近码本条目的承诺损失（Commitment loss）
torch.Tensor[1] – 用于更新码本的码本损失（Codebook loss）
torch.Tensor[B x T] – 码本索引（输入的量化离散表示）
torch.Tensor[B x D x T] – 投影隐向量（量化前的输入连续表示）

embed_code(embed_id: Tensor)[source]

使用码本权重嵌入 ID。

此方法利用码本权重嵌入给定的 ID。

参数:: embed_id (torch.Tensor) – 需要嵌入的包含 ID 的张量。
返回:: 应用码本权重后嵌入的输出张量。
返回类型:: torch.Tensor

decode_code(embed_id: Tensor)[source]

通过转置维度解码嵌入的 ID。

此方法通过对 embed_code 方法的输出张量维度应用转置操作来解码嵌入的 ID。

参数:: embed_id (torch.Tensor) – 包含嵌入 ID 的张量。
返回:: 解码后的张量
返回类型:: torch.Tensor

decode_latents(latents: Tensor)[source]

通过与码本比较，将隐向量表示解码为离散码。

参数:: latents (torch.Tensor) – 待解码的隐向量张量表示。
返回:: 一个包含解码后的隐向量张量 (z_q) 和码索引的元组。
返回类型:: Tuple[torch.Tensor, torch.Tensor]

class speechbrain.lobes.models.discrete.dac.ResidualVectorQuantize(input_dim: int = 512, n_codebooks: int = 9, codebook_size: int = 1024, codebook_dim: int | list = 8, quantizer_dropout: float = 0.0)[source]

基类: Module

在 SoundStream 中引入：一个端到端神经音频编解码器 https://arxiv.org/abs/2107.03312

参数:

input_dim (int, 可选, 默认为 512)
n_codebooks (int, 可选, 默认为 9)
codebook_size (int, 可选, 默认为 1024)
codebook_dim (Union[int, list], 可选, 默认为 8)
quantizer_dropout (float, 可选, 默认为 0.0)

示例

使用预训练的 RVQ 单元。

>>> dac = DAC(load_pretrained=True, model_type="16KHz", model_bitrate="8kbps", tag="latest")
>>> quantizer = dac.quantizer
>>> continuous_embeddings = torch.randn(1, 1024, 20) # Example shape: [Batch, Channels, Time]
>>> discrete_embeddings, codes, _, _, _ = quantizer(continuous_embeddings)

forward(z, n_quantizers: int | None = None)[source]

使用一组固定的 n 个码本对输入张量进行量化，并返回相应的码本向量

参数:

z (torch.Tensor) – 形状 [B x D x T]
n_quantizers (int, 可选) –
要使用的量化器数量 (n_quantizers < self.n_codebooks 例如：用于量化器 dropout)。注意：如果 self.quantizer_dropout 为 True，则忽略此参数

在训练模式下，并使用随机数量的量化器。

返回:

z (torch.Tensor[B x D x T]) – 输入的量化连续表示
codes (torch.Tensor[B x N x T]) – 每个码本的码本索引（输入的量化离散表示）
latents (torch.Tensor[B x N*D x T]) – 投影隐向量（量化前的输入连续表示）
vq/commitment_loss (torch.Tensor[1]) – 用于训练编码器以预测更接近码本条目的承诺损失
vq/codebook_loss (torch.Tensor[1]) – 用于更新码本的码本损失

from_codes(codes: Tensor)[source]

给定量化码，重构连续表示

参数:: codes (torch.Tensor[B x N x T]) – 输入的量化离散表示
返回:: 输入的量化连续表示
返回类型:: torch.Tensor[B x D x T]

from_latents(latents: Tensor)[source]

给定未量化的隐向量，重构量化后的连续表示。

参数:

latents (torch.Tensor[B x N x T]) – 投影后的输入连续表示

返回:

torch.Tensor[B x D x T] – 完全投影空间的量化表示
torch.Tensor[B x D x T] – 隐向量空间的量化表示

class speechbrain.lobes.models.discrete.dac.Snake1d(channels)[source]

基类: Module

实现在 1D 中 Snake 激活函数的 PyTorch 模块。

参数:: channels (int) – 输入张量中的通道数量。

forward(x)[source]

参数:: x (torch.Tensor)
返回类型:: torch.Tensor

class speechbrain.lobes.models.discrete.dac.ResidualUnit(dim: int = 16, dilation: int = 1)[source]

基类: Module

用于卷积神经网络的残差单元模块。

参数:

dim (int, 可选) – 输入张量中的通道数量。默认为 16。
dilation (int, 可选) – 卷积层的扩张率。默认为 1。

forward(x: Tensor) → Tensor[source]

参数:: x (torch.Tensor)
返回类型:: torch.Tensor

class speechbrain.lobes.models.discrete.dac.EncoderBlock(dim: int = 16, stride: int = 1)[source]

基类: Module

用于卷积神经网络的编码器块模块。

此模块构建一个编码器块，由一系列 ResidualUnits 和最终的 Snake1d 激活函数组成，后跟加权归一化 1D 卷积。此块可用作自编码器等架构中的编码器的一部分。

参数:

dim (int, 可选) – 输出通道数量。默认为 16。
stride (int, 可选) – 最终卷积层的步幅。默认为 1。

forward(x: Tensor)[source]

参数:: x (torch.Tensor)
返回类型:: torch.Tensor

class speechbrain.lobes.models.discrete.dac.Encoder(d_model: int = 64, strides: list = [2, 4, 8, 8], d_latent: int = 64)[source]

基类: Module

用于 DAC 编码器部分的 PyTorch 模块。

参数:

d_model (int, 可选) – 模型的初始维度。默认为 64。
strides (list, 可选) – 每个 EncoderBlock 中用于下采样的步幅值列表。默认为 [2, 4, 8, 8]。
d_latent (int, 可选) – 输出隐向量空间的维度。默认为 64。

示例

创建 Encoder 实例 >>> encoder = Encoder() >>> audio_input = torch.randn(1, 1, 16000) # 示例形状：[批量, 通道, 时间] >>> continuous_embedding = encoder(audio_input)

使用预训练的编码器。

>>> dac = DAC(load_pretrained=True, model_type="16KHz", model_bitrate="8kbps", tag="latest")
>>> encoder = dac.encoder
>>> audio_input = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time]
>>> continuous_embeddings = encoder(audio_input)

forward(x)[source]

参数:: x (torch.Tensor)
返回类型:: torch.Tensor

class speechbrain.lobes.models.discrete.dac.DecoderBlock(input_dim: int = 16, output_dim: int = 8, stride: int = 1)[source]

基类: Module

表示解码器架构中一个块的 PyTorch 模块。

参数:

input_dim (int, 可选) – 输入通道数量。默认为 16。
output_dim (int, 可选) – 输出通道数量。默认为 8。
stride (int, 可选) – 转置卷积的步幅，控制上采样。默认为 1。

forward(x)[source]

参数:: x (torch.Tensor)
返回类型:: torch.Tensor

class speechbrain.lobes.models.discrete.dac.Decoder(input_channel: int, channels: int, rates: List[int], d_out: int = 1)[source]

基类: Module

用于 DAC 解码器部分的 PyTorch 模块。

参数:

input_channel (int) – 输入张量中的通道数量。
channels (int) – 卷积层的基本通道数量。
rates (list) – 每个解码器块的步幅率列表
d_out (int) – 最终卷积层的输出维度，默认为 1。

示例

创建 Decoder 实例

>>> decoder = Decoder(128, 256,  [8, 8, 4, 2])
>>> discrete_embeddings = torch.randn(1, 128, 20) # Example shape: [Batch, Channels, Time]
>>> recovered_audio = decoder(discrete_embeddings)

使用预训练的解码器。请注意，实际输入应该是适当的离散表示。这里使用随机生成的输入是为了说明用法。

>>> dac = DAC(load_pretrained=True, model_type="16KHz", model_bitrate="8kbps", tag="latest")
>>> decoder = dac.decoder
>>> discrete_embeddings = torch.randn(1, 1024, 20) # Example shape: [Batch, Channels, Time]
>>> recovered_audio = decoder(discrete_embeddings)

forward(x)[source]

参数:: x (torch.Tensor)
返回类型:: torch.Tensor

class speechbrain.lobes.models.discrete.dac.DAC(encoder_dim: int = 64, encoder_rates: List[int] = [2, 4, 8, 8], latent_dim: int | None = None, decoder_dim: int = 1536, decoder_rates: List[int] = [8, 8, 4, 2], n_codebooks: int = 9, codebook_size: int = 1024, codebook_dim: int | list = 8, quantizer_dropout: bool = False, sample_rate: int = 44100, model_type: str = '44khz', model_bitrate: str = '8kbps', tag: str = 'latest', load_path: str | Path | None = None, strict: bool = False, load_pretrained: bool = False)[source]

基类: Module

用于音频数据编码和解码的离散自编码器编解码器 (DAC)。

此类实现了带有量化功能的自编码器架构，用于高效处理音频。它包括编码器、量化器和解码器，用于将音频数据转换为压缩的隐向量表示并重构回音频。此实现支持初始化新模型和加载预训练模型。

参数:

encoder_dim (int) – 编码器的维度。
encoder_rates (List[int]) – 每个编码器层的下采样率。
latent_dim (int, 可选) – 隐向量空间的维度，如果为 None 则自动计算。
decoder_dim (int) – 解码器的维度。
decoder_rates (List[int]) – 每个解码器层的上采样率。
n_codebooks (int) – 向量量化的码本数量。
codebook_size (int) – 每个码本的大小。
codebook_dim (Union[int, list]) – 每个码本条目的维度。
quantizer_dropout (bool) – 是否在量化器中使用 dropout。
sample_rate (int) – 音频数据的采样率。
model_type (str) – 要加载的模型类型（如果已预训练）。
model_bitrate (str) – 要加载的模型的比特率（如果已预训练）。
tag (str) – 要加载的模型的特定标签（如果已预训练）。
load_path (str, 可选) – 加载预训练模型的路径，如果为 None 则自动下载。
strict (bool) – 是否严格执行状态字典匹配。
load_pretrained (bool) – 是否加载预训练模型。

示例

创建新的 DAC 实例

>>> dac = DAC()
>>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time]
>>> tokens, embeddings = dac(audio_data)

加载预训练的 DAC 实例

>>> dac = DAC(load_pretrained=True, model_type="16KHz", model_bitrate="8kbps", tag="latest")
>>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time]
>>> tokens, embeddings = dac(audio_data)

上面获得或从其他来源获得的 tokens 和离散 embeddings 可以被解码

>>> dac = DAC(load_pretrained=True, model_type="16KHz", model_bitrate="8kbps", tag="latest")
>>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time]
>>> tokens, embeddings = dac(audio_data)
>>> decoded_audio = dac.decode(embeddings)

encode(audio_data: Tensor, n_quantizers: int | None = None)[source]

编码给定的音频数据并返回量化隐向量码

参数:

audio_data (torch.Tensor[B x 1 x T]) – 要编码的音频数据
n_quantizers (int, 可选) – 要使用的量化器数量，默认为 None。如果为 None，则使用所有量化器。

返回:

“z” (torch.Tensor[B x D x T]) – 输入的量化连续表示
”codes” (torch.Tensor[B x N x T]) – 每个码本的码本索引（输入的量化离散表示）
”latents” (torch.Tensor[B x N*D x T]) – 投影隐向量（量化前的输入连续表示）
”vq/commitment_loss” (torch.Tensor[1]) – 用于训练编码器以预测更接近码本条目的承诺损失
”vq/codebook_loss” (torch.Tensor[1]) – 用于更新码本的码本损失
”length” (int) – 输入音频的采样数

decode(z: Tensor)[source]

解码给定的隐向量码并返回音频数据

参数:: z (torch.Tensor) – 形状 [B x D x T] 输入的量化连续表示
返回:: torch.Tensor – 解码后的音频数据。
返回类型:: 形状 B x 1 x length

forward(audio_data: Tensor, sample_rate: int | None = None, n_quantizers: int | None = None)[source]

模型前向传播

参数:

audio_data (torch.Tensor[B x 1 x T]) – 要编码的音频数据
sample_rate (int, 可选) – 音频数据的采样率（Hz），默认为 None。如果为 None，则默认为 self.sample_rate
n_quantizers (int, 可选) – 要使用的量化器数量，默认为 None。如果为 None，则使用所有量化器。

返回:

“tokens” (torch.Tensor[B x N x T]) – 每个码本的码本索引（输入的量化离散表示）
”embeddings” (torch.Tensor[B x D x T]) – 输入的量化连续表示