attention 机制在PyTorch中如何使用

在PyTorch中使用attention机制，通常需要使用nn.MultiheadAttention类，该类实现了多头注意力机制。使用该类的步骤如下：

1. 导入必要的库

import torch
import torch.nn as nn

2. 定义输入

query = torch.randn(10, 32, 512)
key = torch.randn(20, 32, 512)
value = torch.randn(20, 32, 512)

其中，query、key和value都是三维张量，分别表示查询、键和值。第一维表示序列长度，第二维表示batch size，第三维表示特征维度。

3. 实例化MultiheadAttention类

attention = nn.MultiheadAttention(embed_dim=512, num_heads=8)

其中，embed_dim表示特征维度，num_heads表示注意力头数。

4. 进行注意力计算

output, attention_weights = attention(query, key, value)

其中，output表示输出结果，attention_weights表示注意力权重。

完整代码示例：

import torch
import torch.nn as nn

# 定义输入
query = torch.randn(10, 32, 512)
key = torch.randn(20, 32, 512)
value = torch.randn(20, 32, 512)

# 实例化MultiheadAttention类
attention = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# 进行注意力计算
output, attention_weights = attention(query, key, value)

print(output.shape)  # torch.Size([10, 32, 512])
print(attention_weights.shape)  # torch.Size([32, 10, 20])