speechbrain.lobes.models.transformer.Branchformer 模块

Branchformer 实现。

参考：“Branchformer：并行MLP-注意力架构，用于捕捉语音识别和理解中的局部和全局上下文”

来源：代码的某些部分可能改编自ESPNet。

作者 * Titouan Parcollet 2023

摘要

类：

BranchformerEncoder	该类实现了Branchformer编码器。
BranchformerEncoderLayer	这是Branchformer编码器层的实现。
ConvolutionBranch	这是Branchformer中卷积分支的实现。

参考

class speechbrain.lobes.models.transformer.Branchformer.ConvolutionBranch(input_size, linear_units=3072, kernel_size=31, activation=<class 'torch.nn.modules.activation.GELU'>, gate_activation=<class 'torch.nn.modules.linear.Identity'>, dropout=0.0, use_linear_after_conv=False)

基础：Module

这是Branchformer中卷积分支的实现。

默认结构是： LN -> 通道投影 -> GeLU -> (CNN 空间门控) -> 通道投影 -> Dropout

Parameters:

• input_size (int) – 特征（通道）维度的预期大小。

• linear_units (int, optional) – 隐藏线性单元中的神经元数量。

• kernel_size (int, optional) – 非瓶颈卷积层的核大小。

• activation (torch.nn.Module, optional) – 预投影后使用的激活函数。

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

• dropout (float, optional) – 丢弃率。

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

Example

>>> x = torch.rand((8, 60, 512))
>>> net = ConvolutionBranch(512, 1024)
>>> output = net(x)
>>> output.shape
torch.Size([8, 60, 512])

forward(x)

Parameters:

x (torch.Tensor -> (B, T, D))

class speechbrain.lobes.models.transformer.Branchformer.BranchformerEncoderLayer(d_model, nhead, kernel_size=31, kdim=None, vdim=None, activation=<class 'torch.nn.modules.activation.GELU'>, dropout=0.0, attention_type='RelPosMHAXL', csgu_linear_units=3072, gate_activation=<class 'torch.nn.modules.linear.Identity'>, use_linear_after_conv=False)

基础：Module

这是Branchformer编码器层的实现。

Parameters:

• d_model (int) – 输入嵌入的预期大小。

• nhead (int) – 注意力头的数量。

• kernel_size (int, optional) – 卷积模型的核大小。

• kdim (int, optional) – 键的维度。

• vdim (int, optional) – 值的维度。

• activation (torch.nn.Module) – 在每个Conformer层中使用的激活函数。

• dropout (int, optional) – 编码器的丢弃率。

• attention_type (str, 可选) – 注意力层的类型，例如 regularMHA 表示常规的多头注意力机制。

• csgu_linear_units (int, optional) – CSGU模块中隐藏线性单元的神经元数量。

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

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

Example

>>> import torch
>>> x = torch.rand((8, 60, 512))
>>> pos_embs = torch.rand((1, 2*60-1, 512))
>>> net = BranchformerEncoderLayer(nhead=8, d_model=512, kernel_size=3)
>>> output = net(x, pos_embs=pos_embs)
>>> output[0].shape
torch.Size([8, 60, 512])

forward(x, src_mask: Tensor | None = None, src_key_padding_mask: Tensor | None = None, pos_embs: Tensor | None = None)

Parameters:

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

• src_mask (torch.Tensor, optional) – 源序列的掩码。

• src_key_padding_mask (torch.Tensor, optional) – 每个批次的源键的掩码。

• pos_embs (torch.Tensor, torch.nn.Module, optional) – 包含输入序列位置嵌入的模块或张量

class speechbrain.lobes.models.transformer.Branchformer.BranchformerEncoder(num_layers, d_model, nhead, kernel_size=31, kdim=None, vdim=None, activation=<class 'torch.nn.modules.activation.GELU'>, dropout=0.0, attention_type='RelPosMHAXL', csgu_linear_units=3072, gate_activation=<class 'torch.nn.modules.linear.Identity'>, use_linear_after_conv=False, output_hidden_states=False, layerdrop_prob=0.0)

基础：Module

该类实现了Branchformer编码器。

Parameters:

• num_layers (int) – 层数。

• d_model (int) – 嵌入维度大小。

• nhead (int) – 注意力头的数量。

• kernel_size (int, optional) – 卷积模型的核大小。

• kdim (int, optional) – 键的维度。

• vdim (int, optional) – 值的维度。

• activation (torch.nn.Module) – 在每个Confomer层中使用的激活函数。

• dropout (int, optional) – 编码器的丢弃率。

• attention_type (str, 可选) – 注意力层的类型，例如 regularMHA 表示常规的多头注意力机制。

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

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

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

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

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

Example

>>> import torch
>>> x = torch.rand((8, 60, 512))
>>> pos_emb = torch.rand((1, 2*60-1, 512))
>>> net = BranchformerEncoder(1, 512, 8)
>>> output, _ = net(x, pos_embs=pos_emb)
>>> output.shape
torch.Size([8, 60, 512])

>>> import torch
>>> x = torch.rand((8, 60, 512))
>>> pos_emb = torch.rand((1, 2*60-1, 512))
>>> net = BranchformerEncoder(1, 512, 8, output_hidden_states=True)
>>> output, attn_list, hidden_list = net(x, pos_embs=pos_emb)
>>> hidden_list[0].shape
torch.Size([8, 60, 512])
>>> len(hidden_list)
2

forward(src, src_mask: Tensor | None = None, src_key_padding_mask: Tensor | None = None, pos_embs: Tensor | None = None, dynchunktrain_config=None)

Parameters:

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

• src_mask (torch.Tensor, optional) – 源序列的掩码。

• src_key_padding_mask (torch.Tensor, optional) – 每个批次的源键的掩码。

• pos_embs (torch.Tensor, torch.nn.Module,) – 包含输入序列位置嵌入的模块或张量 如果提供了自定义的pos_embs，它需要具有形状 (1, 2*S-1, E) 其中S是序列长度，E是嵌入维度。

• dynchunktrain_config (None) – 此编码器不支持此配置。

Returns:

• output (torch.Tensor) – Conformer的输出。

• attention_lst (list) – 注意力值。

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