多头自注意力机制 pytorch

多头自注意力机制是一种用于处理序列数据的机制，它可以将输入序列中的每个元素与其他元素进行交互，从而获得更好的表示。在PyTorch中，可以使用`nn.MultiheadAttention`模块来实现多头自注意力机制。该模块接受三个输入：查询（query）、键（key）和值（value），并输出注意力加权的值。

具体来说，`nn.MultiheadAttention`模块将查询、键和值分别通过线性变换映射到不同的空间中，然后将它们分成多个头（head），每个头都进行注意力计算，最后将多个头的结果拼接起来并通过另一个线性变换得到最终输出。

以下是一个使用`nn.MultiheadAttention`模块实现多头自注意力机制的示例代码：

import torch
import torch.nn as nn

batch_size = 16
seq_len = 10
input_size = 32
num_heads = 4
hidden_size = 64

# 构造输入
x = torch.randn(batch_size, seq_len, input_size)
# 定义多头自注意力机制模块
self_attn = nn.MultiheadAttention(hidden_size, num_heads)

# 将输入通过线性变换映射到hidden_size维空间中
query = nn.Linear(input_size, hidden_size)(x)
key = nn.Linear(input_size, hidden_size)(x)
value = nn.Linear(input_size, hidden_size)(x)

# 计算多头自注意力
output, attn_weights = self_attn(query, key, value)

# 输出结果
print(output.shape)  # torch.Size([batch_size, seq_len, hidden_size])
print(attn_weights.shape)  # torch.Size([batch_size, num_heads, seq_len, seq_len])

在上面的代码中，我们首先构造了一个大小为`(batch_size, seq_len, input_size)`的输入张量`x`，然后定义了一个包含4个头、每个头的隐藏层大小为64的多头自注意力机制模块`self_attn`。接着，我们将输入通过线性变换映射到隐藏层空间中，并将映射后的结果作为查询、键和值输入到`self_attn`模块中进行多头自注意力计算。最后，我们得到了输出张量`output`和注意力权重张量`attn_weights`。