pytorch如何添加一个注意力层

要在PyTorch中添加一个注意力层，可以使用 torch.nn.MultiheadAttention() 函数。此函数能够计算输入序列的特征权重并将其应用于输出序列上，从而实现注意力机制。以下是一个简单的例子：

import torch.nn as nn

# 定义输入和输出维度
input_dim = 512
output_dim = 256

# 创建一个MultiheadAttention层
attention = nn.MultiheadAttention(input_dim, num_heads=8)

# 输入数据x的形状为(batch_size, seq_len, input_dim)
x = torch.randn(32, 100, input_dim)

# 通过MultiheadAttention计算输出y
y = attention(x, x, x)

# 输出y的形状为(batch_size, seq_len, output_dim)
print(y.shape)

在上面的代码中，我们创建了一个具有8个头的MultiheadAttention层，并将其应用于一个具有100个时间步长和512维输入的批次大小为32的张量。注意，MultiheadAttention()函数接受三个输入：查询、键和数值。在这个例子中，我们将它们都设置为相同的张量x。最后，我们得到了一个具有256维输出的张量y。