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

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

Author

舒巴姆·古普塔 2023

摘要

类：

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

函数：

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

参考

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

将权重归一化应用于一维卷积层。

Parameters:

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

Returns:

权重归一化的 nn.Conv1d 层。

Return type:

torch.nn.Module

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

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

Parameters:

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

Returns:

权重归一化的 nn.ConvTranspose1d 层。

Return type:

torch.nn.Module

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

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

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

Parameters:

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

Returns:

模型保存的本地路径。

Return type:

路径

Raises:

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

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

基础：Module

向量量化的实现

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

Factorized codes: Perform nearest neighbor lookup in low-dimensional space
为了提高码本的使用

l2-normalized codes: Converts euclidean distance to cosine similarity which
提高训练稳定性

Parameters:

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

forward(z: Tensor)[source]

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

Parameters:

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

Returns:

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

embed_code(embed_id: Tensor)[source]

使用码本权重嵌入ID。

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

Parameters:: embed_id (torch.Tensor) – 包含需要嵌入的ID的张量。
Returns:: 应用码本权重后的嵌入输出张量。
Return type:: torch.Tensor

decode_code(embed_id: Tensor)[source]

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

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

Parameters:: embed_id (torch.Tensor) – 包含嵌入ID的张量。
Returns:: 解码后的张量
Return type:: torch.Tensor

decode_latents(latents: Tensor)[source]

通过将潜在表示与代码本进行比较，将其解码为离散代码。

Parameters:: latents (torch.Tensor) – 要解码的潜在张量表示。
Returns:: 一个包含解码后的潜在张量（z_q）和代码索引的元组。
Return type:: 元组[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

Parameters:

input_dim (int, 可选, 默认为 512)
n_codebooks (int, 可选, 默认为 9)
codebook_size (int, 可选, 默认为 1024)
codebook_dim (联合[int, 列表], 可选, 默认为 8)
quantizer_dropout (float, 可选, 默认为 0.0)

Example

使用预训练的RVQ单元。

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

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

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

Parameters:

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

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

Returns:

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]

给定量化代码，重建连续表示

Parameters:: codes (torch.Tensor[B x N x T]) – 输入的量化离散表示
Returns:: 输入的量化连续表示
Return type:: torch.Tensor[B x D x T]

from_latents(latents: Tensor)[source]

给定未量化的潜在变量，在量化后重建连续表示。

Parameters:

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

Returns:

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

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

基础：Module

一个实现一维Snake激活函数的PyTorch模块。

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

forward(x)[source]

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

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

基础：Module

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

Parameters:

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

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

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

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

基础：Module

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

该模块构建了一个由一系列ResidualUnits和一个最终的Snake1d激活函数组成的编码器块，随后是一个加权归一化的一维卷积。该块可以用作自动编码器等架构中编码器的一部分。

Parameters:

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

forward(x: Tensor)[source]

Parameters:: x (torch.Tensor)
Return type:: 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模块。

Parameters:

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

Example

创建一个编码器实例 >>> encoder = Encoder() >>> audio_input = torch.randn(1, 1, 44100) # 示例形状：[批次, 通道, 时间] >>> continuous_embedding = encoder(audio_input)

使用预训练的编码器。

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

forward(x)[source]

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

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

基础：Module

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

Parameters:

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

forward(x)[source]

Parameters:: x (torch.Tensor)
Return type:: 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模块。

Parameters:

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

Example

创建一个解码器实例

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

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

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

forward(x)[source]

Parameters:: x (torch.Tensor)
Return type:: 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, 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 = None, strict: bool = False, load_pretrained: bool = False)[source]

基础：Module

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

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

Parameters:

encoder_dim (int) – 编码器的维度。
encoder_rates (List[int]) – 每个编码器层的下采样率。
latent_dim (int, optional) – 潜在空间的维度，如果为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) – 要加载的模型的比特率（如果预训练）。
标签 (str) – 要加载的模型的特定标签（如果预训练）。
load_path (str, 可选) – 加载预训练模型的路径，如果为None则自动下载。
strict (bool) – 是否严格强制执行状态字典匹配。
load_pretrained (bool) – 是否加载预训练模型。

Example

创建一个新的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="44KHz", model_bitrate="8kbps", tag="latest")
>>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time]
>>> tokens, embeddings = dac(audio_data)

上述或其他来源获得的标记和离散嵌入可以被解码：

>>> dac = DAC(load_pretrained=True, model_type="44KHz", 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)[source]

对给定的音频数据进行编码并返回量化后的潜在代码

Parameters:

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

Returns:

“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]

解码给定的潜在代码并返回音频数据

Parameters:: z (torch.Tensor) – 形状 [B x D x T] 输入的量化连续表示
Returns:: torch.Tensor – 解码后的音频数据。
Return type:: 形状 B x 1 x 长度

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

模型前向传播

Parameters:

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

Returns:

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