SE注意力机制代码

下面是一个简单的自注意力机制（self-attention）的代码示例：

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, hidden_size):
        super(SelfAttention, self).__init__()
        self.hidden_size = hidden_size
        self.query = nn.Linear(hidden_size, hidden_size)
        self.key = nn.Linear(hidden_size, hidden_size)
        self.value = nn.Linear(hidden_size, hidden_size)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        query = self.query(x)
        key = self.key(x)
        value = self.value(x)

        # 计算 scaled dot-product attention
        scores = torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.hidden_size))
        attn_weights = self.softmax(scores)

        # 乘上 attention weights
        context = torch.matmul(attn_weights, value)

        return context

这个类中包含了三个线性层：query、key、value，它们将输入的 x 映射到隐藏空间中，然后计算注意力得分。计算得分后，应用 softmax 函数来获得归一化的注意力权重，并使用这些权重对 value 进行加权求和。最后返回加权和作为上下文向量。

可以将这个模块嵌入到模型的任何一层中，以增强模型的表达能力。