如何使用nn.MultiheadAttention

nn.MultiheadAttention是用于多头自注意力机制的PyTorch模块。它可以被用于各种任务，例如自然语言处理和计算机视觉。

以下是使用nn.MultiheadAttention的步骤：

1. 导入必要的库

import torch
import torch.nn as nn

2. 定义输入

# 输入有三个维度：序列长度、批次大小、特征向量的维度
seq_length = 10
batch_size = 32
input_dim = 64

# 随机生成输入张量
inputs = torch.randn(seq_length, batch_size, input_dim)

3. 定义MultiheadAttention模块

# 定义MultiheadAttention模块
attention = nn.MultiheadAttention(embed_dim=input_dim, num_heads=8)

这个模块有两个参数：embed_dim是输入张量的维度，num_heads是要使用的头的数量。

4. 计算自注意力

# 计算自注意力
outputs, attention_weights = attention(inputs, inputs, inputs)

这里我们将输入张量传递了三次，因为MultiheadAttention需要三个输入：查询、键和值。在这种情况下，我们使用相同的输入作为查询、键和值。

输出是两个张量：outputs和attention_weights。outputs是注意力计算后的张量，attention_weights是注意力权重张量，它表示了每个序列元素对于其他元素的重要性。

5. 查看输出

print("outputs shape:", outputs.shape)
print("attention_weights shape:", attention_weights.shape)

输出应该类似于：

outputs shape: torch.Size([10, 32, 64])
attention_weights shape: torch.Size([32, 10, 10])

这意味着注意力计算后的张量的形状是[seq_length, batch_size, input_dim]，而注意力权重张量的形状是[batch_size, seq_length, seq_length]。