pytorch给神经网络添加注意力层

在PyTorch中，可以使用`torch.nn.MultiheadAttention`模块来添加注意力层。以下是一个简单的示例代码：

import torch
import torch.nn as nn

class AttentionLayer(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(AttentionLayer, self).__init__()
        self.multihead_attn = nn.MultiheadAttention(hidden_size, num_heads=1)
        self.linear = nn.Linear(input_size, hidden_size)
    
    def forward(self, x):
        # x shape: (seq_len, batch_size, input_size)
        query = self.linear(x)
        out, _ = self.multihead_attn(query, x, x)
        return out

在上面的代码中，`AttentionLayer`类的初始化方法中创建了一个`nn.MultiheadAttention`对象和一个线性层。`nn.MultiheadAttention`对象需要传入一个隐藏大小和头的数量，这里我们将头的数量设置为1。在`forward`方法中，首先将输入`x`通过线性层进行变换得到查询向量`query`，然后使用`nn.MultiheadAttention`进行注意力计算，最后返回注意力计算结果。

需要注意的是，`nn.MultiheadAttention`模块接受的输入形状为(seq_len, batch_size, hidden_size)，因此需要将输入进行转置以满足要求。