speechbrain.lobes.models.beats 模块

此 lobe 支持集成预训练的 BEATs: 使用声学分词器进行音频预训练。

参考: https://arxiv.org/abs/2212.09058 基于 Github 源码: https://github.com/microsoft/unilm/tree/master/beats

您可以从以下链接下载检查点: https://github.com/microsoft/unilm/tree/master/beats

作者
  • Pooneh Mousavi 2024

摘要

BEATs

BEATs: 使用声学分词器进行音频预训练。

BEATsConfig

BEATs 模型的配置类。

GLU_Linear

实现门控线性单元 (GLU) 与线性变换的组合。

GradMultiply

一个在反向传播过程中缩放梯度的自定义 autograd 函数。

MultiheadAttention

实现多头注意力,支持相对位置嵌入和门控相对位置嵌入(基于 GRU)等高级特性。

SamePad

实现一个模块,用于在卷积后调整张量的填充以保持其原始大小,并提供因果填充选项。

Swish

将 Swish 激活函数实现为一个 PyTorch 模块。

TransformerEncoder

实现 Transformer Encoder 模块。

TransformerSentenceEncoderLayer

实现单个 Transformer Sentence Encoder 层。

函数

gelu

应用高斯误差线性单元 (GELU) 激活函数。

gelu_accurate

使用精确近似应用高斯误差线性单元 (GELU) 激活函数。

get_activation_fn

返回与提供的激活名称对应的激活函数。

init_bert_params

初始化 BERT 模型中模块的权重和偏置。

quant_noise

包装模块并对其权重应用量化噪声,以便后续使用迭代乘积量化 (iPQ) 进行量化。

参考

class speechbrain.lobes.models.beats.BEATs(ckp_path: str = None, freeze: bool = True, output_all_hiddens: bool = False)[source]

基类: Module

BEATs: 使用声学分词器进行音频预训练。

此类实现了 BEATs 模型,该模型处理音频信号用于特征提取或下游任务。该模型支持从检查点加载、应用归一化以及可选地冻结参数。

参数:
  • ckp_path (str, 可选) – 检查点文件的路径。如果为 None,模型将不使用预训练权重进行初始化。您可以从以下链接下载检查点: https://github.com/microsoft/unilm/tree/master/beats

  • freeze (bool, 可选 (默认: False)) – 如果为 True,模型参数将被冻结,并且模型设置为评估模式。

  • output_all_hiddens (bool, 可选 (默认: False)) – 如果为 True,forward 函数将输出所有 transformer 层的隐藏状态。例如,BEATs_iter3 有 12 个 transformer 层,输出形状为 (13, B, T, C),其中 CNN 输出的投影被添加到开头。如果为 False,forward 函数仅输出最后一个 transformer 层的隐藏状态。

示例

>>> audio = torch.randn(4, 10000)  # Batch of 4 audio signals
>>> length = torch.tensor([1.0, 0.5, 0.75, 1.0])
>>> model = BEATs()
>>> outputs = model.extract_features(audio, length)[0]
>>> outputs.shape
torch.Size([4, 24, 768])
forward_padding_mask(features: Tensor, padding_mask: Tensor) Tensor[source]

调整给定特征的填充掩码。

参数:
  • features (torch.Tensor) – 补丁嵌入后的输入特征。

  • padding_mask (torch.Tensor) – 输入信号的原始填充掩码。

返回:

调整后的填充掩码。

返回类型:

torch.Tensor

preprocess(source: Tensor, fbank_mean: float = 15.41663, fbank_std: float = 6.55582) Tensor[source]

通过提取滤波器组并应用归一化来预处理输入波形。

参数:
  • source (torch.Tensor) – 输入波形信号。

  • fbank_mean (float, 可选) – 用于滤波器组归一化的均值(默认值:15.41663)。

  • fbank_std (float, 可选) – 用于滤波器组归一化的标准差(默认值:6.55582)。

返回:

归一化的滤波器组。

返回类型:

torch.Tensor

forward(wav: Tensor, wav_lens: Tensor | None = None, fbank_mean: float = 15.41663, fbank_std: float = 6.55582)[source]

接收输入波形并返回其对应的 beats 编码。

参数:
  • wav (torch.Tensor) – 要转换为特征的音频信号批量。

  • wav_lens (torch.Tensor) – 以 SpeechBrain 格式给定的 wav 的相对长度。

  • fbank_mean (float, 可选) – 用于滤波器组归一化的均值(默认值:15.41663)。

  • fbank_std (float, 可选) – 用于滤波器组归一化的标准差(默认值:6.55582)。

返回类型:

BEATs 编码特征。

extract_features(wav: Tensor, wav_lens: Tensor | None = None, fbank_mean: float = 15.41663, fbank_std: float = 6.55582) Tensor[source]

从输入波形中提取特征。

参数:
  • wav (torch.Tensor) – 要转换为特征的音频信号批量。

  • wav_lens (torch.Tensor) – 以 SpeechBrain 格式给定的 wav 的相对长度。

  • fbank_mean (float, 可选) – 用于滤波器组归一化的均值(默认值:15.41663)。

  • fbank_std (float, 可选) – 用于滤波器组归一化的标准差(默认值:6.55582)。

返回:

从 BEATs 模型中提取的特征。

返回类型:

torch.Tensor

speechbrain.lobes.models.beats.gelu_accurate(x)[source]

使用精确近似应用高斯误差线性单元 (GELU) 激活函数。

参数:

x (torch.Tensor) – 要应用 GELU 激活的输入张量。

返回:

应用了 GELU 激活的逐元素张量。

返回类型:

torch.Tensor

speechbrain.lobes.models.beats.gelu(x: Tensor) Tensor[source]

应用高斯误差线性单元 (GELU) 激活函数。

参数:

x (torch.Tensor) – 要应用 GELU 激活的输入张量。

返回:

应用了 GELU 激活的逐元素张量。

返回类型:

torch.Tensor

speechbrain.lobes.models.beats.get_activation_fn(activation: str)[source]

返回与提供的激活名称对应的激活函数。

参数:

activation (str) – 激活函数名称。支持的值: - “relu”: 应用 ReLU 激活。- “gelu”: 应用 GELU 激活。- “gelu_fast”: 带弃用警告的 gelu_accurate 别名。- “gelu_accurate”: 应用精确的 GELU 激活。- “tanh”: 应用 Tanh 激活。- “linear”: 应用恒等函数。- “glu”: 应用恒等函数(GLU 占位符)。

返回:

应用于输入张量的相应激活函数。

返回类型:

Callable[[torch.Tensor], torch.Tensor]

抛出:

RuntimeError – 如果指定的激活函数不受支持。

class speechbrain.lobes.models.beats.SamePad(kernel_size, causal=False)[source]

基类: Module

实现一个模块,用于在卷积后调整张量的填充以保持其原始大小,并提供因果填充选项。

这对于处理卷积层中的填充特别有用,其中核大小或因果性会影响输出大小。

参数:
  • kernel_size (int) – 卷积核的大小。

  • causal (bool, 可选 (默认值=False)) – 如果为 True,则通过从张量末尾移除 (kernel_size - 1) 个元素来应用因果填充。如果为 False,则移除元素以中心对齐填充,确保输出大小与输入大小匹配。

forward(x)[source]

调整输入张量 x 的填充。

如果 self.remove > 0,该方法将根据 kernel_sizecausal 设置沿最后一个维度对张量进行切片,以移除多余的填充。

参数:

x (torch.Tensor) – 要调整填充的输入张量。

返回:

调整填充后的张量。

返回类型:

torch.Tensor

class speechbrain.lobes.models.beats.Swish[source]

基类: Module

将 Swish 激活函数实现为一个 PyTorch 模块。

Swish 是一个平滑的非单调激活函数,定义如下:

Swish(x) = x * sigmoid(x)

在深度学习中,它常用于改善某些架构的训练性能。

forward(x)[source]

将 Swish 激活函数应用于输入张量。

参数:

x (torch.Tensor) – 应用 Swish 激活的输入张量。

返回:

应用 Swish 激活后的输入张量。

返回类型:

torch.Tensor

class speechbrain.lobes.models.beats.GLU_Linear(input_dim, output_dim, glu_type='sigmoid', bias_in_glu=True)[source]

基类: Module

实现门控线性单元 (GLU) 与线性变换的组合。

参数:
  • input_dim (int) – 输入特征的维度。

  • output_dim (int) – 输出特征的维度。

  • glu_type (str, optional (default="sigmoid")) – 用于门控的激活函数类型。支持的值包括:- “sigmoid”:使用 sigmoid 激活函数。- “swish”:使用 Swish 激活函数。- “relu”:使用 ReLU 激活函数。- “gelu”:使用 GELU 激活函数。

  • bias_in_glu (bool, optional (default=True)) – 是否在线性变换中包含偏置项。

class speechbrain.lobes.models.beats.GradMultiply(*args, **kwargs)[source]

Bases: Function

一个在反向传播过程中缩放梯度的自定义 autograd 函数。

这在需要在不影响前向传播输出的情况下缩放梯度的情况下非常有用。前向传播原样返回输入,而后向传播则按指定的因子缩放梯度。

static forward(ctx, x, scale)[source]

执行 GradMultiply 函数的前向传播。

参数:
  • ctx (torch.autograd.Function) – 用于存储后向计算信息的上下文对象。

  • x (torch.Tensor) – 将原样转发的输入张量。

  • scale (float) – 后向传播期间梯度将按此因子缩放。

返回:

与输入张量相同的新张量。

返回类型:

torch.Tensor

static backward(ctx, grad)[source]

执行后向传播,按存储的因子缩放梯度。

参数:
  • ctx (torch.autograd.Function) – 包含存储缩放因子的上下文对象。

  • grad (torch.Tensor) – 来自后续层的梯度张量。

返回:

缩放后的梯度张量和 None(对于 scale 输入,它没有梯度)。

返回类型:

Tuple[torch.Tensor, None]

speechbrain.lobes.models.beats.quant_noise(module, p, block_size)[source]

包装模块并对其权重应用量化噪声,以便后续使用迭代乘积量化 (iPQ) 进行量化。

此方法在论文《Training with Quantization Noise for Extreme Model Compression》中有所描述。它在训练期间引入量化噪声,以提高模型在极端权重压缩场景下的鲁棒性。

参数:
  • module (nn.Module) – 将应用量化噪声的模块。支持的模块包括 Linear、Embedding 和 Conv2d。

  • p (float) – 应用的量化噪声量。通常是一个概率或缩放因子。

  • block_size (int) – 用于后续 iPQ 量化的块大小。

返回类型:

None

class speechbrain.lobes.models.beats.TransformerEncoder(args)[source]

基类: Module

实现 Transformer Encoder 模块。

参数:

args (Namespace or dict) – 模型超参数和配置的集合。

forward(x, padding_mask=None, output_all_hiddens=None)[source]

通过 Transformer Encoder 层处理输入序列。

参数:
  • x (torch.Tensor) – 输入张量,形状为 (seq_len, batch_size, embed_dim),包含输入嵌入。

  • padding_mask (torch.Tensor, optional) – 二进制掩码,形状为 (batch_size, seq_len),指示哪些位置是填充且在注意力计算中应被忽略。默认为 None

  • output_all_hiddens (bool, optional) – 如果为 True,则除了最终输出外,还会返回所有 encoder 层的隐藏状态。默认为 None

返回:

  • 最终输出张量,形状为 (seq_len, batch_size, embed_dim)

返回类型:

Tuple[torch.Tensor, List[torch.Tensor]]

extract_features(x, padding_mask=None, output_all_hiddens=None)[source]

通过位置卷积、层归一化、dropout 和一系列 Transformer Encoder 层从输入序列中提取特征。

参数:
  • x (torch.Tensor) – 输入张量,形状为 (batch_size, seq_len, embed_dim),包含输入嵌入。

  • padding_mask (torch.Tensor, optional) – 二进制掩码,形状为 (batch_size, seq_len),指示哪些位置是填充且在计算中应被忽略。默认为 None

  • output_all_hiddens (bool, optional) – 如果为 True,则除了最终输出外,还会收集并返回所有 encoder 层的隐藏状态。默认为 None

返回:

  • 最终输出张量,形状为 (batch_size, seq_len, embed_dim)

返回类型:

Tuple[torch.Tensor, List[torch.Tensor]]

class speechbrain.lobes.models.beats.TransformerSentenceEncoderLayer(embedding_dim: float = 768, ffn_embedding_dim: float = 3072, num_attention_heads: float = 8, dropout: float = 0.1, attention_dropout: float = 0.1, activation_dropout: float = 0.1, activation_fn: str = 'relu', layer_norm_first: bool = False, deep_norm: bool = False, has_relative_attention_bias: bool = False, num_buckets: int = 0, max_distance: int = 0, rescale_init: bool = False, gru_rel_pos: bool = False, encoder_layers: int = 0)[source]

基类: Module

实现单个 Transformer Sentence Encoder 层。

参数:
  • embedding_dim (float, optional (default=768)) – 输入嵌入的维度。

  • ffn_embedding_dim (float, optional (default=3072)) – 前馈网络隐藏层的维度。

  • num_attention_heads (float, optional (default=8)) – 自注意力的注意力头数。

  • dropout (float, optional (default=0.1)) – 应用于前馈网络和注意力层输出的 dropout 比率。

  • attention_dropout (float, optional (default=0.1)) – 应用于注意力机制内部的 dropout 比率。

  • activation_dropout (float, optional (default=0.1)) – 应用于前馈网络中激活函数后的 dropout 比率。

  • activation_fn (str, optional (default="relu")) – 前馈网络中使用的激活函数。支持的值包括“relu”和“gelu”。

  • layer_norm_first (bool, optional (default=False)) – 如果为 True,则在注意力层和前馈层之前应用层归一化;否则,在之后应用。

  • deep_norm (bool, optional (default=False)) – 如果为 True,则对残差连接使用深度归一化缩放。

  • has_relative_attention_bias (bool, optional (default=False)) – 如果为 True,则在注意力机制中包含相对位置偏置。

  • num_buckets (int, optional (default=0)) – 用于相对注意力偏置(如果启用)的桶数。

  • max_distance (int, optional (default=0)) – 用于相对注意力偏置(如果启用)的最大距离。

  • rescale_init (bool, optional (default=False)) – 如果为 True,则重新缩放参数初始化以提高稳定性。

  • gru_rel_pos (bool, optional (default=False)) – 如果为 True,则包含 GRU 风格的相对位置编码。

  • encoder_layers (int, optional (default=0)) – Transformer 中的 encoder 层数。

forward(x: Tensor, self_attn_mask: Tensor = None, self_attn_padding_mask: Tensor = None, need_weights: bool = False, pos_bias=None)[source]

通过 Transformer sentence encoder 层处理输入张量。

参数:
  • x (torch.Tensor) – 输入张量,形状为 (seq_len, batch_size, embed_dim)

  • self_attn_mask (torch.Tensor, optional) – 自注意力机制的掩码,通常用于因果或填充掩码。默认为 None

  • self_attn_padding_mask (torch.Tensor, optional) – 填充掩码,形状为 (batch_size, seq_len),指示在注意力计算中应忽略哪些 token。默认为 None

  • need_weights (bool, optional (default=False)) – 是否返回注意力权重。如果为 True,则输出中包含注意力权重。

  • pos_bias (optional) – 相对注意力的位置偏置,如果适用。默认为 None

返回:

  • x (torch.Tensor): 应用 encoder 层后

的输出张量,形状为 (seq_len, batch_size, embed_dim)

返回类型:

Tuple[torch.Tensor, torch.Tensor, optional]

class speechbrain.lobes.models.beats.MultiheadAttention(embed_dim, num_heads, kdim=None, vdim=None, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, self_attention=False, encoder_decoder_attention=False, q_noise=0.0, qn_block_size=8, has_relative_attention_bias=False, num_buckets=32, max_distance=128, gru_rel_pos=False, rescale_init=False)[source]

基类: Module

实现多头注意力,支持相对位置嵌入和门控相对位置嵌入(基于 GRU)等高级特性。

参数:
  • embed_dim (int) – 输入嵌入的总维度。

  • num_heads (int) – 注意力头数。

  • kdim (int, optional) – 键嵌入的维度。默认为 embed_dim

  • vdim (int, optional) – 值嵌入的维度。默认为 embed_dim

  • dropout (float, optional) – 注意力权重的 dropout 概率。默认为 0.0。

  • bias (bool, optional) – 是否在投影中包含偏置项。默认为 True。

  • add_bias_kv (bool, optional) – 是否为键和值投影包含偏置项。默认为 False。

  • add_zero_attn (bool, optional) – 是否包含零注意力向量。默认为 False。

  • self_attention (bool, optional) – 该层是否用于自注意力。默认为 False。

  • encoder_decoder_attention (bool, optional) – 该层是否用于 encoder-decoder 注意力。默认为 False。

  • q_noise (float, optional) – 量化噪声水平。默认为 0.0。

  • qn_block_size (int, optional) – 量化的块大小。默认为 8。

  • has_relative_attention_bias (bool, optional) – 是否使用相对位置嵌入。默认为 False。

  • num_buckets (int, optional) – 相对位置嵌入的桶数。默认为 32。

  • max_distance (int, optional) – 相对位置嵌入的最大距离。默认为 128。

  • gru_rel_pos (bool, optional) – 是否使用门控相对位置嵌入。默认为 False。

  • rescale_init (bool, optional) – 是否重新缩放权重初始化。默认为 False。

reset_parameters()[source]

初始化投影层和相对位置嵌入的权重。

compute_bias(query_length: int, key_length: int) Tensor[source]

计算注意力分数的相对位置偏置。

参数:
  • query_length (int) – 查询序列的长度。

  • key_length (int) – 键序列的长度。

返回:

一个张量,形状为 (num_heads, query_length, key_length),包含每个注意力头的相对位置偏置值。

返回类型:

torch.Tensor

forward(query: Tensor, key: Tensor | None, value: Tensor | None = None, key_padding_mask: Tensor | None = None, incremental_state: Dict[str, Dict[str, Tensor | None]] | None = None, need_weights: bool = True, static_kv: bool = False, attn_mask: Tensor | None = None, before_softmax: bool = False, need_head_weights: bool = False, position_bias: Tensor | None = None) Tuple[Tensor, Tensor | None, Tensor | None][source]

执行多头注意力的前向传播,支持相对位置嵌入、缓存和可选的 dropout。

此方法实现了多头注意力的核心功能,并支持相对位置偏置、增量解码和各种掩码选项等可选功能。

参数:
  • query (torch.Tensor) – 查询张量,形状为 (target_length, batch_size, embed_dim)

  • key (torch.Tensor, optional) – 键张量,形状为 (source_length, batch_size, embed_dim)。默认为 None

  • value (torch.Tensor, optional) – 值张量,形状为 (source_length, batch_size, embed_dim)。默认为 None

  • key_padding_mask (torch.Tensor, optional) – 用于排除填充键的掩码,形状为 (batch_size, source_length),填充元素用 1 表示。默认为 None

  • incremental_state (dict, optional) – 存储用于增量解码的缓存键和值张量。默认为 None

  • need_weights (bool, optional) – 如果为 True,则返回注意力权重。默认为 True

  • static_kv (bool, optional) – 如果为 True,则在增量解码时,键和值张量保持静态。默认为 False

  • attn_mask (torch.Tensor, optional) – 注意力掩码,用于阻止某些位置参与注意力计算,通常用于因果注意力。形状:(target_length, source_length)。默认为 None

  • before_softmax (bool, optional) – 如果为 True,则返回 softmax 之前的原始注意力分数。默认为 False

  • need_head_weights (bool, optional) – 如果为 True,则返回每个注意力头的权重。这意味着 need_weights=True。默认为 False

  • position_bias (torch.Tensor, optional) – 预计算的位置偏置张量。如果为 None,则在前向传播期间计算。

返回:

  • attn (torch.Tensor) – 注意力输出张量,形状为 (target_length, batch_size, embed_dim)

  • attn_weights (torch.Tensor, optional) – 注意力权重张量,形状为 (batch_size, num_heads, target_length, source_length),如果 need_head_weights=False,则跨头平均。

  • position_bias (torch.Tensor, optional) – 计算或传入的相对位置偏置,形状为 (num_heads, target_length, source_length)

apply_bias(k, v, bsz, attn_mask=None, key_padding_mask=None)[source]

将 bias_k 和 bias_v 应用于键和值张量,并相应地更新注意力掩码和键填充掩码。

参数:
  • k (torch.Tensor) – 键张量。

  • v (torch.Tensor) – 值张量。

  • bsz (int) – 批量大小。

  • attn_mask (torch.Tensor) – 注意力掩码

  • key_padding_mask (torch.Tensor) – 键填充掩码。

返回:

Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] – 注意力掩码和键填充掩码。

返回类型:

更新后的键、值、

apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int)[source]

对注意力权重应用稀疏掩码。

参数:
  • attn_weights (torch.Tensor) – 注意力权重张量,形状为 (batch_size * num_heads, tgt_len, src_len)

  • tgt_len (int) – 目标序列长度。

  • src_len (int) – 源序列长度。

  • bsz (int) – 批量大小。

返回:

(可能被修改的)注意力权重张量。默认情况下,这与输入张量相同。

返回类型:

torch.Tensor

speechbrain.lobes.models.beats.init_bert_params(module: Module) None[source]

初始化 BERT 模型中模块的权重和偏置。

参数:

module (nn.Module) – 要初始化的模块。可以是 nn.Linearnn.EmbeddingMultiheadAttention 中的一个。

class speechbrain.lobes.models.beats.BEATsConfig(cfg=None)[source]

Bases: object

BEATs 模型的配置类。

此类定义了 BEATs 模型的配置。它提供了一个可以通过 update 方法使用自定义设置更新的默认配置。

参数:

cfg (dict, optional) – 包含自定义配置值的字典。如果提供,它将覆盖默认设置。

update(cfg: dict)[source]

使用给定配置字典中的键值对更新实例的属性。

参数:

cfg (dict) – 包含用于更新实例的配置值的字典。