cross-attention代码pytorch

以下是一个简单的 PyTorch 实现交叉注意力（cross-attention）的代码：

import torch
import torch.nn as nn

class CrossAttention(nn.Module):
    def __init__(self, input_dim):
        super(CrossAttention, self).__init__()
        self.input_dim = input_dim
        self.linear_q = nn.Linear(input_dim, input_dim)
        self.linear_k = nn.Linear(input_dim, input_dim)
        self.linear_v = nn.Linear(input_dim, input_dim)
        
    def forward(self, x1, x2):
        q = self.linear_q(x1)
        k = self.linear_k(x2)
        v = self.linear_v(x2)
        
        attn_weights = torch.matmul(q, k.transpose(1, 2))
        attn_weights = torch.softmax(attn_weights, dim=-1)
        
        output = torch.matmul(attn_weights, v)
        
        return output

其中，`input_dim` 表示输入的特征维度，`x1` 和 `x2` 分别表示两个输入张量。在 `forward` 方法中，我们首先通过三个线性层将输入张量映射到相同的特征空间，然后计算注意力权重并将其应用于第二个输入张量 `x2` 的值上，最后返回加权和。