speechbrain.lobes.models.transformer.TransformerASR 模块

用于ASR的Transformer，采用SpeechBrain风格。

作者 * Jianyuan Zhong 2020 * Titouan Parcollet 2024 * Luca Della Libera 2024

摘要

类：

EncoderWrapper

这是任何ASR变压器编码器的包装器。

TransformerASR

这是用于ASR的transformer模型的实现。

TransformerASRStreamingContext

用于TransformerASR实例的流式元数据和状态。

函数：

make_transformer_src_mask

准备源变换器掩码，该掩码根据因果或其他简单的受限注意力方法限制哪些帧可以关注哪些帧。

make_transformer_src_tgt_masks

此函数生成用于训练transformer模型的掩码，针对ASR上下文进行了定制，包括编码掩码，以及可选的解码掩码（如果指定了tgt）。

参考

class speechbrain.lobes.models.transformer.TransformerASR.TransformerASRStreamingContext(dynchunktrain_config: DynChunkTrainConfig, encoder_context: Any)[source]

基础类：object

流式元数据和状态用于一个TransformerASR实例。

dynchunktrain_config: DynChunkTrainConfig

动态块训练配置，包含块大小和上下文大小信息。

encoder_context: Any

不透明的编码器上下文信息。它由编码器的make_streaming_context方法构造，并在使用encode_streaming时传递给编码器。

speechbrain.lobes.models.transformer.TransformerASR.make_transformer_src_mask(src: Tensor, causal: bool = False, dynchunktrain_config: DynChunkTrainConfig | None = None) → Tensor | None[source]

准备源变换器掩码，该掩码根据因果或其他简单的受限注意力方法限制哪些帧可以关注哪些帧。

Parameters:

  • src (torch.Tensor) – 用于构建掩码的源张量。张量的内容目前并未实际使用；仅使用其形状和其他元数据（例如设备）。

  • 因果性 (bool) – 是否应使用严格的因果性。帧将无法关注任何未来的帧。

  • dynchunktrain_config (DynChunkTrainConfig, 可选) – 动态块训练配置。这实现了一种简单的块状注意力形式。与causal不兼容。

Returns:

一个形状为 (timesteps, timesteps) 的布尔掩码张量。

Return type:

torch.Tensor

speechbrain.lobes.models.transformer.TransformerASR.make_transformer_src_tgt_masks(src, tgt=None, wav_len=None, pad_idx=0, causal: bool = False, dynchunktrain_config: DynChunkTrainConfig | None = None)[source]

此函数生成用于训练transformer模型的掩码， 针对ASR上下文进行了定制，带有编码掩码，并且可选地带有解码掩码（如果指定了tgt）。

Parameters:

  • src (torch.Tensor) – 编码器的输入序列（必需）。

  • tgt (torch.Tensor) – 解码器的输入序列。

  • wav_len (torch.Tensor) – 输入的长度。

  • pad_idx (int) – <pad> 标记的索引（默认=0）。

  • 因果 (布尔) – 是否应使用严格的因果关系。参见 make_asr_src_mask

  • dynchunktrain_config (DynChunkTrainConfig, 可选) – 动态块训练配置。参见 make_asr_src_mask

Returns:

  • src_key_padding_mask (torch.Tensor) – 用于忽略填充的键填充掩码

  • tgt_key_padding_mask (torch.Tensor) – 用于忽略填充的键填充掩码

  • src_mask (torch.Tensor) – 用于忽略无效（例如未来）时间步的掩码

  • tgt_mask (torch.Tensor) – 用于忽略无效（例如未来）时间步的掩码

class speechbrain.lobes.models.transformer.TransformerASR.TransformerASR(tgt_vocab, input_size, d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, d_ffn=2048, dropout=0.1, activation=<class 'torch.nn.modules.activation.ReLU'>, positional_encoding='fixed_abs_sine', normalize_before=False, kernel_size: int | None = 31, bias: bool | None = True, encoder_module: str | None = 'transformer', conformer_activation: ~torch.nn.modules.module.Module | None = <class 'speechbrain.nnet.activations.Swish'>, branchformer_activation: ~torch.nn.modules.module.Module | None = <class 'torch.nn.modules.activation.GELU'>, attention_type: str | None = 'regularMHA', max_length: int | None = 2500, causal: bool | None = None, csgu_linear_units: int | None = 3072, gate_activation: ~torch.nn.modules.module.Module | None = <class 'torch.nn.modules.linear.Identity'>, use_linear_after_conv: bool | None = False, output_hidden_states=False, layerdrop_prob=0.0)[source]

基础类: TransformerInterface

这是ASR的transformer模型的实现。

该架构基于论文“Attention Is All You Need”： https://arxiv.org/pdf/1706.03762.pdf

Parameters:

  • tgt_vocab (int) – 词汇表的大小。

  • input_size (int) – 输入特征大小。

  • d_model (int, optional) – 嵌入维度大小。 (默认=512).

  • nhead (int, 可选) – 多头注意力模型中的头数（默认=8）。

  • num_encoder_layers (int, 可选) – 编码器中子编码层的数量（默认=6）。

  • num_decoder_layers (int, 可选) – 解码器中子解码层的数量（默认=6）。

  • d_ffn (int, 可选) – 前馈网络模型的维度（默认=2048）。

  • dropout (int, 可选) – 丢弃率的值（默认=0.1）。

  • activation (torch.nn.Module, optional) – FFN层的激活函数。 推荐使用：relu 或 gelu（默认=relu）。

  • positional_encoding (str, 可选) – 使用的位置编码类型。例如，‘fixed_abs_sine’ 表示固定的绝对位置编码。

  • normalize_before (bool, optional) – 是否在Transformer层中的MHA或FFN之前或之后应用归一化。 默认为True，因为这被证明可以带来更好的性能和训练稳定性。

  • kernel_size (int, optional) – 当使用Conformer时，卷积层中的核大小。

  • bias (bool, 可选) – 是否在Conformer卷积层中使用偏置。

  • encoder_module (str, 可选) – 在编码器中选择Conformer或Transformer。解码器固定为Transformer。

  • conformer_activation (torch.nn.Module, optional) – 在Conformer卷积层之后使用的激活模块。例如Swish、ReLU等。它必须是一个torch模块。

  • branchformer_activation (torch.nn.Module, optional) – 用于Branchformer编码器中的激活模块。例如Swish、ReLU等。它必须是一个torch模块。

  • attention_type (str, 可选) – 在所有Transformer或Conformer层中使用的注意力层类型。 例如 regularMHA 或 RelPosMHA。

  • max_length (int, 可选) – 输入中目标和源序列的最大长度。 用于位置编码。

  • causal (bool, 可选) – 编码器是否应该是因果的（解码器总是因果的）。 如果是因果的，Conformer卷积层是因果的。

  • csgu_linear_units (int, 可选) – CSGU模块中隐藏线性单元的神经元数量。 -> Branchformer

  • gate_activation (torch.nn.Module, optional) – 用于CSGU模块门控的激活函数。 -> Branchformer

  • use_linear_after_conv (bool, 可选) – 如果为True，将应用一个大小为input_size//2的线性变换。 -> Branchformer

  • output_hidden_states (bool, 可选) – 模型是否应该输出隐藏状态作为张量列表。

  • layerdrop_prob (float) – 丢弃整个层的概率。

Example

>>> src = torch.rand([8, 120, 512])
>>> tgt = torch.randint(0, 720, [8, 120])
>>> net = TransformerASR(
...     720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU
... )
>>> enc_out, dec_out = net.forward(src, tgt)
>>> enc_out.shape
torch.Size([8, 120, 512])
>>> dec_out.shape
torch.Size([8, 120, 512])
forward(src, tgt, wav_len=None, pad_idx=0)[source]
Parameters:

  • src (torch.Tensor) – 编码器的输入序列。

  • tgt (torch.Tensor) – 解码器的输入序列。

  • wav_len (torch.Tensor, optional) – 形状为 (batch, ) 的 Torch Tensor，包含每个示例的相对长度与填充长度的比例。

  • pad_idx (int, 可选) – 标记的索引（默认=0）。

Returns:

  • encoder_out (torch.Tensor) – 编码器的输出。

  • decoder_out (torch.Tensor) – 解码器的输出

  • hidden_state_lst (list, optional) – 编码器隐藏层的输出。 仅在output_hidden_states设置为true时有效。

decode(tgt, encoder_out, enc_len=None)[source]

该方法实现了变压器模型的解码步骤。

Parameters:

  • tgt (torch.Tensor) – 解码器的输入序列。

  • encoder_out (torch.Tensor) – 编码器的隐藏输出。

  • enc_len (torch.LongTensor) – 编码器状态的实际长度。

Return type:

预测

encode(src, wav_len=None, pad_idx=0, dynchunktrain_config: DynChunkTrainConfig | None = None)[source]

编码器前向传递

Parameters:

  • src (torch.Tensor) – 编码器的输入序列。

  • wav_len (torch.Tensor, optional) – 形状为 (batch, ) 的 Torch Tensor，包含每个示例的相对长度与填充长度的比例。

  • pad_idx (int) – 用于填充的索引。

  • dynchunktrain_config (DynChunkTrainConfig) – 动态分块配置。

Returns:

encoder_out

Return type:

torch.Tensor

encode_streaming(src, context: TransformerASRStreamingContext)[source]

流式编码器前向传递

Parameters:

  • src (torch.Tensor) – 编码器的输入序列（块）。

  • 上下文 (TransformerASRStreamingContext) – 对流式上下文的可变引用。这保存了在多个块推理之间需要持久化的状态，并且可以使用make_streaming_context来构建。这个上下文会被此函数修改。

Return type:

此块的编码器输出。

Example

>>> import torch
>>> from speechbrain.lobes.models.transformer.TransformerASR import TransformerASR
>>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
>>> net = TransformerASR(
...     tgt_vocab=100,
...     input_size=64,
...     d_model=64,
...     nhead=8,
...     num_encoder_layers=1,
...     num_decoder_layers=0,
...     d_ffn=128,
...     attention_type="RelPosMHAXL",
...     positional_encoding=None,
...     encoder_module="conformer",
...     normalize_before=True,
...     causal=False,
... )
>>> ctx = net.make_streaming_context(DynChunkTrainConfig(16, 1))
>>> src1 = torch.rand([8, 16, 64])
>>> src2 = torch.rand([8, 16, 64])
>>> out1 = net.encode_streaming(src1, ctx)
>>> out1.shape
torch.Size([8, 16, 64])
>>> ctx.encoder_context.layers[0].mha_left_context.shape
torch.Size([8, 16, 64])
>>> out2 = net.encode_streaming(src2, ctx)
>>> out2.shape
torch.Size([8, 16, 64])
>>> ctx.encoder_context.layers[0].mha_left_context.shape
torch.Size([8, 16, 64])
>>> combined_out = torch.concat((out1, out2), dim=1)
>>> combined_out.shape
torch.Size([8, 32, 64])
make_streaming_context(dynchunktrain_config: DynChunkTrainConfig, encoder_kwargs={})[source]

为此转换器及其编码器创建一个空白的流式上下文。

Parameters:

  • dynchunktrain_config (DynChunkTrainConfig) – 运行时分块注意力配置。

  • encoder_kwargs (dict) – 传递给编码器的make_streaming_context的参数。 编码器所需的元数据可能因编码器而异。

Return type:

TransformerASRStreamingContext

class speechbrain.lobes.models.transformer.TransformerASR.EncoderWrapper(transformer, *args, **kwargs)[source]

基础：Module

这是任何ASR变换器编码器的包装器。默认情况下，TransformerASR的.forward()函数会进行编码和解码。使用此包装器后，.forward()函数仅变为.encode()。

重要：TransformerASR 类必须包含一个 .encode() 函数。

Parameters:

  • transformer (sb.lobes.models.TransformerInterface) – 一个包含 .encode() 函数的 Transformer 实例。

  • *args (元组)

  • **kwargs (dict) – 传递给父类的参数。

Example

>>> src = torch.rand([8, 120, 512])
>>> tgt = torch.randint(0, 720, [8, 120])
>>> net = TransformerASR(
...     720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU
... )
>>> encoder = EncoderWrapper(net)
>>> enc_out = encoder(src)
>>> enc_out.shape
torch.Size([8, 120, 512])
forward(x, wav_lens=None, pad_idx=0, **kwargs)[source]

处理输入张量 x 并返回输出张量。

forward_streaming(x, context)[source]

处理输入的音频块张量 x，使用并更新可变的编码器 context

make_streaming_context(*args, **kwargs)[source]

初始化一个流式上下文。将所有参数转发给底层的transformer。参见speechbrain.lobes.models.transformer.TransformerASR.make_streaming_context()。