speechbrain.nnet.attention 模块

实现注意力模块的库。

作者

Ju-Chieh Chou 2020
Jianyuan Zhong 2020
Loren Lugosch 2020
Samuele Cornell 2020
Shucong Zhang 2024

摘要

类

`ContentBasedAttention`	此类实现了用于 seq2seq 学习的基于内容的注意力模块。
`KeyValueAttention`	此类实现了用于 seq2seq 学习的单头键值注意力模块。
`LocationAwareAttention`	此类实现了用于 seq2seq 学习的位置感知注意力模块。
`MemoiseAtLeastSize`	记忆一个函数，该函数将其第一个参数作为调用底层函数的最小值指示。
`MultiheadAttention`	此类是 torch.nn.MultiHeadAttention 的多头注意力包装器。
`PositionalwiseFeedForward`	此类实现了“Attention Is All You Need”中的位置前馈模块。
`PrecomputedRoPESinusoids`	用于缓存旋转位置嵌入 (RoPE) 所需的正弦和余弦，以便旋转向量。
`RelPosEncXL`	RelPosMHAXL 的相对位置编码。
`RelPosMHAXL`	此类实现了类似于 Transformer XL https://arxiv.org/pdf/1901.02860.pdf 中的相对多头实现。
`RoPEMHA`	这是使用 RoPE 位置嵌入的多头自注意力实现。

函数

`masks_union`	这是一个实用函数，用于将 SpeechBrain 中的标准 key_padding_mask 和 attn_mask 组合成一个用于 scaled_dot_product_attention。此函数不支持对 attn_score 进行加权。因此，如果希望使用浮点值作为掩码，则不应使用此函数。
`memoise_at_least`	装饰器，记忆一个函数，该函数将其第一个参数作为调用底层函数的最小值指示。如果记忆中存储了匹配先前函数调用的结果，则将返回存储的结果，而不是再次调用该函数。

参考

class speechbrain.nnet.attention.ContentBasedAttention(enc_dim, dec_dim, attn_dim, output_dim, scaling=1.0)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

此类实现了用于 seq2seq 学习的基于内容的注意力模块。

参考资料：通过联合学习对齐和翻译实现神经机器翻译, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

参数:

enc_dim (int) – 编码器层的大小。
dec_dim (int) – 解码器层的大小。
attn_dim (int) – 注意力特征的大小。
output_dim (int) – 输出上下文向量的大小。
scaling (float) – 控制锐化程度的因子（默认值：1.0）。

示例

>>> enc_tensor = torch.rand([4, 10, 20])
>>> enc_len = torch.ones([4]) * 10
>>> dec_tensor = torch.rand([4, 25])
>>> net = ContentBasedAttention(enc_dim=20, dec_dim=25, attn_dim=30, output_dim=5)
>>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
>>> out_tensor.shape
torch.Size([4, 5])

reset()[源代码]: 重置注意力模块中的内存。

forward(enc_states, enc_len, dec_states)[源代码]

返回注意力模块的输出。

参数:

enc_states (torch.Tensor) – 待注意的张量。
enc_len (torch.Tensor) – 每句话的真实长度（不含填充）。
dec_states (torch.Tensor) – 查询张量。

返回类型:

注意力模块的输出。

class speechbrain.nnet.attention.LocationAwareAttention(enc_dim, dec_dim, attn_dim, output_dim, conv_channels, kernel_size, scaling=1.0)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

此类实现了用于 seq2seq 学习的位置感知注意力模块。

参考：Attention-Based Models for Speech Recognition, Chorowski et.al. https://arxiv.org/pdf/1506.07503.pdf

参数:

enc_dim (int) – 编码器的大小。
dec_dim (int) – 解码器的大小。
attn_dim (int) – 注意力特征的大小。
output_dim (int) – 输出上下文向量的大小。
conv_channels (int) – 位置特征的通道数。
kernel_size (int) – 位置特征卷积层的大小。
scaling (float) – 控制锐化程度的因子（默认值：1.0）。

示例

>>> enc_tensor = torch.rand([4, 10, 20])
>>> enc_len = torch.ones([4]) * 10
>>> dec_tensor = torch.rand([4, 25])
>>> net = LocationAwareAttention(
...     enc_dim=20,
...     dec_dim=25,
...     attn_dim=30,
...     output_dim=5,
...     conv_channels=10,
...     kernel_size=100)
>>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
>>> out_tensor.shape
torch.Size([4, 5])

precomputed_enc_h: Tensor | None

reset()[源代码]: 重置注意力模块中的内存。

forward(enc_states, enc_len, dec_states)[源代码]

返回注意力模块的输出。

参数:

enc_states (torch.Tensor) – 待注意的张量。
enc_len (torch.Tensor) – 每句话的真实长度（不含填充）。
dec_states (torch.Tensor) – 查询张量。

返回类型:

注意力模块的输出。

class speechbrain.nnet.attention.KeyValueAttention(enc_dim, dec_dim, attn_dim, output_dim)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

此类实现了用于 seq2seq 学习的单头键值注意力模块。

参考：“Attention Is All You Need”，Vaswani 等人，第 3.2.1 节

参数:

enc_dim (int) – 计算键和值的编码器特征向量的大小。
dec_dim (int) – 计算查询的解码器特征向量的大小。
attn_dim (int) – 注意力特征的大小。
output_dim (int) – 输出上下文向量的大小。

示例

>>> enc_tensor = torch.rand([4, 10, 20])
>>> enc_len = torch.ones([4]) * 10
>>> dec_tensor = torch.rand([4, 25])
>>> net = KeyValueAttention(enc_dim=20, dec_dim=25, attn_dim=30, output_dim=5)
>>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
>>> out_tensor.shape
torch.Size([4, 5])

reset()[源代码]: 重置注意力模块中的内存。

forward(enc_states, enc_len, dec_states)[源代码]

返回注意力模块的输出。

参数:

enc_states (torch.Tensor) – 待注意的张量。
enc_len (torch.Tensor) – 每句话的真实长度（不含填充）。
dec_states (torch.Tensor) – 查询张量。

返回类型:

注意力模块的输出。

class speechbrain.nnet.attention.RelPosEncXL(emb_dim: int, dtype: dtype = torch.float32)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

RelPosMHAXL 的相对位置编码。

参数:

emb_dim (int) – 嵌入的大小，也控制位置嵌入最后一个维度的大小
dtype (torch.dtype, optional) – 如果未指定，默认为 torch.float32。控制输出嵌入的数据类型（但不影响计算精度，计算精度仍为 torch.float32）。

make_pe(seq_len: int)[源代码]

为给定序列长度构建位置嵌入张量。

参数:: seq_len (int) – 要创建位置嵌入的序列长度。
返回:: 形状为 [1, 2*seq_len-1, embed_dim] 的位置嵌入张量
返回类型:: torch.Tensor

forward(x: Tensor)[源代码]

构建位置嵌入张量。类似于 make_pe()，但使用提供的张量的形状信息。

参数:: x (torch.Tensor) – 输入张量，形状为 batch_size, seq_len, embed_dim
返回:: pos_emb – 形状为 [1, 2*seq_len-1, embed_dim] 的位置嵌入张量
返回类型:: torch.Tensor

class speechbrain.nnet.attention.RelPosMHAXL(embed_dim, num_heads, dropout=0.0, vbias=False, vdim=None, mask_pos_future=False)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

此类实现了类似于 Transformer XL https://arxiv.org/pdf/1901.02860.pdf 中的相对多头实现。

参数:

embed_dim (int) – 计算键和值的编码器特征向量的大小。
num_heads (int) – 注意力头的数量。
dropout (float, optional) – Dropout 比率。
vbias (bool, optional) – 是否使用偏差计算值。
vdim (int, optional) – 值的大小。默认为 embed_dim（注意每个头是 embed_dim // num_heads）。
mask_pos_future (bool, optional) – 是否掩盖未来位置编码值。对于因果应用（例如解码器）必须为 True。

示例

>>> inputs = torch.rand([6, 60, 512])
>>> pos_emb = torch.rand([1, 2*60-1, 512])
>>> net = RelPosMHAXL(num_heads=8, embed_dim=inputs.shape[-1])
>>> outputs, attn = net(inputs, inputs, inputs, pos_emb)
>>> outputs.shape
torch.Size([6, 60, 512])

rel_shift(x)[源代码]: 相对位移实现。

forward(query, key, value, pos_embs, key_padding_mask=None, attn_mask=None, return_attn_weights=True)[源代码]

计算注意力。

参数:

query (torch.Tensor) – (B, L, E)，其中 L 是目标序列长度，B 是批量大小，E 是嵌入维度。
key (torch.Tensor) – (B, S, E)，其中 S 是源序列长度，B 是批量大小，E 是嵌入维度。
value (torch.Tensor) – (B, S, E)，其中 S 是源序列长度，B 是批量大小，E 是嵌入维度。
pos_embs (torch.Tensor) – 双向正弦位置嵌入张量 (1, 2*S-1, E)，其中 S 是源序列和目标序列长度之间的最大长度，E 是嵌入维度。
key_padding_mask (torch.Tensor) – (B, S)，其中 B 是批量大小，S 是源序列长度。如果提供了 ByteTensor，则非零位置将被忽略，零位置将保持不变。如果提供了 BoolTensor，则值为 True 的位置将被忽略，值为 False 的位置将保持不变。
attn_mask (torch.Tensor) – 2D 掩码 (L, S)，其中 L 是目标序列长度，S 是源序列长度。3D 掩码 (N*num_heads, L, S)，其中 N 是批量大小，L 是目标序列长度，S 是源序列长度。attn_mask 确保位置 i 允许注意未被掩码的位置。如果提供了 ByteTensor，则非零位置不允许注意，零位置将保持不变。如果提供了 BoolTensor，则值为 True 的位置不允许注意，值为 False 的位置将保持不变。如果提供了 FloatTensor，它将被添加到注意力权重中。
return_attn_weights (bool) – 是否额外返回注意力权重。

返回:

out (torch.Tensor) – (B, L, E)，其中 L 是目标序列长度，B 是批量大小，E 是嵌入维度。
attn_score (torch.Tensor) – (B, L, S)，其中 B 是批量大小，L 是目标序列长度，S 是源序列长度。

class speechbrain.nnet.attention.MultiheadAttention(nhead, d_model, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

此类是 torch.nn.MultiHeadAttention 的多头注意力包装器。

参考：https://pytorch.ac.cn/docs/stable/nn.html

参数:

nhead (int) – 并行注意力头。
d_model (int) – 模型层的大小。
dropout (float) – 应用于 attn_output_weights 的 Dropout 层（默认值：0.0）。
bias (bool) – 添加偏差作为模块参数（默认值：True）。
add_bias_kv (bool) – 在 dim=0 的键和值序列中添加偏差。
add_zero_attn (bool) – 在 dim=1 的键和值序列中添加新的零批量。
kdim (int) – 键中的特征总数（默认值：None）。
vdim (int) – 值中的特征总数（默认值：None）。

示例

>>> inputs = torch.rand([8, 60, 512])
>>> net = MultiheadAttention(nhead=8, d_model=inputs.shape[-1])
>>> outputs, attn = net(inputs, inputs, inputs)
>>> outputs.shape
torch.Size([8, 60, 512])

forward(query, key, value, attn_mask: Tensor | None = None, key_padding_mask: Tensor | None = None, return_attn_weights: bool = True, pos_embs: Tensor | None = None)[源代码]

计算注意力。

参数:

query (torch.Tensor) – (B, L, E)，其中 L 是目标序列长度，B 是批量大小，E 是嵌入维度。
key (torch.Tensor) – (B, S, E)，其中 S 是源序列长度，B 是批量大小，E 是嵌入维度。
value (torch.Tensor) – (B, S, E)，其中 S 是源序列长度，B 是批量大小，E 是嵌入维度。
attn_mask (torch.Tensor, optional) – 2D 掩码 (L, S)，其中 L 是目标序列长度，S 是源序列长度。3D 掩码 (N*num_heads, L, S)，其中 N 是批量大小，L 是目标序列长度，S 是源序列长度。attn_mask 确保位置 i 允许注意未被掩码的位置。如果提供了 ByteTensor，则非零位置不允许注意，零位置将保持不变。如果提供了 BoolTensor，则值为 True 的位置不允许注意，值为 False 的位置将保持不变。如果提供了 FloatTensor，它将被添加到注意力权重中。
key_padding_mask (torch.Tensor, optional) – (B, S)，其中 B 是批量大小，S 是源序列长度。如果提供了 ByteTensor，则非零位置将被忽略，零位置将保持不变。如果提供了 BoolTensor，则值为 True 的位置将被忽略，值为 False 的位置将保持不变。
return_attn_weights (bool, optional) – True 表示额外返回注意力权重，False 表示不返回。
pos_embs (torch.Tensor, optional) – 添加到注意力图中的位置嵌入，形状为 (L, S, E) 或 (L, S, 1)。

返回:

attn_output (torch.Tensor) – (B, L, E)，其中 L 是目标序列长度，B 是批量大小，E 是嵌入维度。
attn_output_weights (torch.Tensor) – (B, L, S)，其中 B 是批量大小，L 是目标序列长度，S 是源序列长度。仅当 return_attn_weights=True 时返回（默认为 True）。

class speechbrain.nnet.attention.PositionalwiseFeedForward(d_ffn, input_shape=None, input_size=None, dropout=0.0, activation: type = <class 'torch.nn.modules.activation.ReLU'>)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

此类实现了“Attention Is All You Need”中的位置前馈模块。

参数:

d_ffn (int) – 隐藏层大小。
input_shape (tuple, optional) – 期望的输入形状。或者使用 input_size。
input_size (int, optional) – 期望的输入大小。或者使用 input_shape。
dropout (float, optional) – Dropout 比率。
activation (torch.nn.Module, optional) – 应用的激活函数（推荐：ReLU, GELU）。

示例

>>> inputs = torch.rand([8, 60, 512])
>>> net = PositionalwiseFeedForward(256, input_size=inputs.shape[-1])
>>> outputs = net(inputs)
>>> outputs.shape
torch.Size([8, 60, 512])

forward(x)[源代码]: 将 PositionalwiseFeedForward 应用于输入张量 x。

class speechbrain.nnet.attention.PrecomputedRoPESinusoids(max_length: int, input_size: int, dtype: dtype, device: device)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

用于缓存旋转位置嵌入 (RoPE) 所需的正弦和余弦，以便旋转向量。这存储了来自 https://arxiv.org/pdf/2104.09864 中公式 (15) 的非零项。

参数:

max_length (int) – 允许的最大输入序列长度。对于其他参数的固定设置，计算时间为 O(max_length)。
input_size (int) – 输入序列中每个向量的大小，即每个注意力头的维度。
dtype (torch.dtype) – 张量的数据类型。
device (torch.device) – 放置张量的 Torch 设备。

示例

>>> precomputed = PrecomputedRoPESinusoids(3, 8, torch.float32, torch.device('cpu'))
>>> precomputed.cosines.shape
torch.Size([3, 8])
>>> precomputed.sines.shape == precomputed.cosines.shape
True
>>> precomputed.cosines
tensor([[ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
        [ 0.5403,  0.5403,  0.9950,  0.9950,  0.9999,  0.9999,  1.0000,  1.0000],
        [-0.4161, -0.4161,  0.9801,  0.9801,  0.9998,  0.9998,  1.0000,  1.0000]])
>>> precomputed.sines
tensor([[-0.0000,  0.0000, -0.0000,  0.0000, -0.0000,  0.0000, -0.0000,  0.0000],
        [-0.8415,  0.8415, -0.0998,  0.0998, -0.0100,  0.0100, -0.0010,  0.0010],
        [-0.9093,  0.9093, -0.1987,  0.1987, -0.0200,  0.0200, -0.0020,  0.0020]])
>>> precomputed.index_swap
tensor([1, 0, 3, 2, 5, 4, 7, 6])

class speechbrain.nnet.attention.MemoiseAtLeastSize(function: Callable, round_up: Callable[[Any], Any])[源代码]

基类: object

记忆一个函数，该函数将其第一个参数作为调用底层函数的最小值指示。

参数:

function (Callable) – 要调用的函数。
round_up (Callable[[Any], Any]) – 一个进行向上取整的函数。此函数取整到的值越少，函数被重复调用的可能性越小。

speechbrain.nnet.attention.memoise_at_least(round_up: Callable[[Any], Any]) → Callable[[Callable], MemoiseAtLeastSize][源代码]

装饰器，记忆一个函数，该函数将其第一个参数作为调用底层函数的最小值指示。如果记忆中存储了匹配先前函数调用的结果，则将返回存储的结果，而不是再次调用该函数。

参数:: round_up (Callable[[Any], Any]) – 一个进行向上取整的函数。这将使用传入的第一个参数调用。底层函数将接收取整后的版本，而不是此第一个参数。此函数取整到的值越少，函数被重复调用的可能性越小。
返回类型:: 传入的函数，但具有 MemoiseAtLeastSize 能力。

class speechbrain.nnet.attention.RoPEMHA(embed_dim, num_heads, dropout=0.0, vbias=False, vdim=None)[源代码]

参考：NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

这是使用 RoPE 位置嵌入的多头自注意力实现。由于它依赖于 Torch 进行自注意力计算，因此比 RelPosMHAXL 快得多，同时提供相同或更好的精度水平。

关于 RoPE 的详细信息：https://arxiv.org/pdf/2104.09864。

参数:

embed_dim (int) – 计算键和值的编码器特征向量的大小。
num_heads (int) – 注意力头的数量。
dropout (float, optional) – Dropout 比率。
vbias (bool, optional) – 是否使用偏差计算值。
vdim (int, optional) – 值的大小。默认为 embed_dim（注意每个头是 embed_dim // num_heads）。

示例

>>> max_len = 64
>>> inputs = torch.rand([6, 60, 512])
>>> num_heads = 8
>>> net = RoPEMHA(num_heads=num_heads, embed_dim=inputs.shape[-1])
>>> outputs, attn = net(inputs, inputs, inputs)
>>> outputs.shape
torch.Size([6, 60, 512])

forward(query, key, value, key_padding_mask=None, attn_mask=None, pos_embs=None, return_attn_weights=True)[源代码]

通过 Pytorch 注意力计算注意力。

参数:

query (torch.Tensor) – (B, L, E)，其中 L 是目标序列长度，B 是批量大小，E 是嵌入维度。
key (torch.Tensor) – (B, S, E)，其中 S 是源序列长度，B 是批量大小，E 是嵌入维度。
value (torch.Tensor) – (B, S, E)，其中 S 是源序列长度，B 是批量大小，E 是嵌入维度。
key_padding_mask (torch.Tensor) – (B, S)，其中 B 是批量大小，S 是源序列长度。如果提供了 ByteTensor，则非零位置将被忽略，零位置将保持不变。如果提供了 BoolTensor，则值为 True 的位置将被忽略，值为 False 的位置将保持不变。
attn_mask (torch.BoolTensor) – 2D 掩码 (L, S)，其中 L 是目标序列长度，S 是源序列长度。值为 True 的位置将被忽略，值为 False 的位置将保持不变。
pos_embs (torch.Tensor) – 此类未使用。保留以保持兼容性。
return_attn_weights (bool) – 是否额外返回注意力权重。

返回:

out (torch.Tensor) – (B, L, E)，其中 L 是目标序列长度，B 是批量大小，E 是嵌入维度。
attn_score (torch.Tensor) – (B, L, S)，其中 B 是批量大小，L 是目标序列长度，S 是源序列长度。

speechbrain.nnet.attention.masks_union(bsz, klen, num_heads, attn_mask, key_padding_mask)[源代码]

这是一个实用函数，用于将 SpeechBrain 中的标准 key_padding_mask 和 attn_mask 组合成一个用于 scaled_dot_product_attention。此函数不支持对 attn_score 进行加权。因此，如果希望使用浮点值作为掩码，则不应使用此函数。

参数:

bsz (int) – 批量大小维度。
klen (int) – 键张量的时间维度（序列长度）。
num_heads (int) – 使用这些掩码的注意力模块的头数。
attn_mask (torch.BoolTensor) – 2D 掩码 (L, S)，其中 L 是目标序列长度，S 是源序列长度。值为 True 的位置将被忽略，值为 False 的位置将保持不变。
key_padding_mask (torch.BoolTensor) – (B, S)，其中 B 是批量大小，S 是源序列长度。值为 True 的位置将被忽略，值为 False 的位置将保持不变。

返回:

out – (bsz, num_heads, klen, klen)，其中 False 值被掩码，True 值未被掩码（与输入张量相反）。

返回类型:

torch.BoolTensor