多头注意力 pytorch

多头注意力是一种在深度学习中常用的注意力机制，它在PyTorch中可以通过torch.nn.MultiheadAttention模块来实现。多头注意力允许模型同时关注输入的不同部分，从而提高模型的表达能力和性能。

在使用torch.nn.MultiheadAttention时，我们需要指定一些参数来定义多头注意力的行为。其中，最重要的参数包括：
- embed_dim：输入的特征维度。
- num_heads：多头注意力的头数，即并行计算的注意力头的个数。
- dropout：用于控制注意力权重的dropout比例。

下面是一个使用torch.nn.MultiheadAttention的示例代码：

import torch
import torch.nn as nn

# 定义输入
batch_size = 16
seq_len = 10
embed_dim = 256
num_heads = 8

# 创建MultiheadAttention实例
multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)

# 生成输入数据
input = torch.randn(seq_len, batch_size, embed_dim)

# 进行多头注意力计算
output, _ = multihead_attn(input, input, input)

# 输出结果
print(output.shape)  # 输出：torch.Size([seq_len, batch_size, embed_dim])

在上述示例中，我们首先创建了一个MultiheadAttention实例，然后生成了输入数据input，最后调用multihead_attn函数进行多头注意力计算。输出的结果是一个三维张量，形状为[seq_len, batch_size, embed_dim]，表示多头注意力的输出特征。