交叉注意力机制的·代码

以下是一个使用PyTorch实现的交叉注意力机制的代码示例：

import torch
import torch.nn as nn

class CrossAttention(nn.Module):
    def __init__(self, d_model):
        super(CrossAttention, self).__init__()
        self.query_linear = nn.Linear(d_model, d_model)
        self.key_linear = nn.Linear(d_model, d_model)
        self.value_linear = nn.Linear(d_model, d_model)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x1, x2):
        query = self.query_linear(x1)
        key = self.key_linear(x2).transpose(-2, -1)
        value = self.value_linear(x2).transpose(-2, -1)
        scores = torch.matmul(query, key)
        distribution = self.softmax(scores)
        out = torch.matmul(distribution, value).transpose(-2, -1)
        return out

这是一个简单的交叉注意力模块，输入是两个特征向量x1和x2，输出是一个新的特征向量out。在模块中，我们首先使用线性层将输入向量映射到一个新的表示空间中，然后计算注意力分布并将其应用于输入特征向量的值。最后，我们使用另一个线性层将注意力加权的值转换回原始表示空间中。