软注意力机制的代码

以下是一个简单的软注意力机制的代码实现，其中假设有一个输入张量input和一个参考张量ref，均为二维张量，注意力权重为一个二维张量，输出为加权的输入张量。

import torch
import torch.nn.functional as F

class SoftAttention(torch.nn.Module):
    def __init__(self, input_dim, ref_dim):
        super(SoftAttention, self).__init__()
        self.input_dim = input_dim
        self.ref_dim = ref_dim
        self.linear = torch.nn.Linear(input_dim, ref_dim)
    
    def forward(self, input, ref):
        # 计算注意力权重
        attn_weights = F.softmax(torch.bmm(self.linear(input), ref.transpose(1, 2)), dim=2)
        # 加权求和
        weighted_input = torch.bmm(attn_weights, input)
        return weighted_input

在上面的代码中，我们首先定义了一个SoftAttention类，它接受两个参数：输入张量的维度input_dim和参考张量的维度ref_dim，它们分别对应输入张量input和参考张量ref的最后一个维度大小。在模型的初始化函数中，我们定义了一个线性层，将输入张量的维度转换为参考张量的维度，用于计算注意力权重。

在forward函数中，我们首先使用torch.bmm函数计算输入张量和线性层的乘积，然后将结果与参考张量的转置的乘积相乘，得到注意力权重。最后，我们使用torch.bmm函数将注意力权重与输入张量相乘，得到加权的输入张量。最后输出加权的输入张量。