pytorch 十字交叉注意力机制代码

十字交叉注意力机制（Cross-Cross Attention Mechanism）是一种用于自然语言处理中的注意力机制。它可以将两个不同位置的序列进行交叉关注，从而实现更好的信息交流和融合。下面是一个基于PyTorch实现的十字交叉注意力机制的代码示例：

import torch
import torch.nn as nn

class CrossAttention(nn.Module):
    def __init__(self, hidden_size):
        super(CrossAttention, self).__init__()
        self.hidden_size = hidden_size
        self.attention = nn.Linear(hidden_size * 2, 1)

    def forward(self, source, target):
        # source: batch_size x source_len x hidden_size
        # target: batch_size x target_len x hidden_size
        batch_size, source_len, _ = source.size()
        _, target_len, _ = target.size()

        # compute attention scores
        source = source.unsqueeze(2).repeat(1, 1, target_len, 1)  # batch_size x source_len x target_len x hidden_size
        target = target.unsqueeze(1).repeat(1, source_len, 1, 1)  # batch_size x source_len x target_len x hidden_size
        combined = torch.cat([source, target], dim=-1)  # batch_size x source_len x target_len x hidden_size*2
        scores = self.attention(combined).squeeze(-1)  # batch_size x source_len x target_len

        # compute context vectors
        attn = nn.functional.softmax(scores, dim=-1)  # batch_size x source_len x target_len
        target = torch.bmm(attn.transpose(1, 2), source)  # batch_size x target_len x hidden_size
        source = torch.bmm(attn, target)  # batch_size x source_len x hidden_size

        return source, target

在这个代码中，我们定义了一个名为`CrossAttention`的`nn.Module`，它接受两个输入源`source`和`target`。这两个源都是三维张量，分别表示形状为`(batch_size, seq_len, hidden_size)`的输入序列。在前向传递过程中，我们首先计算`source`和`target`之间的注意力得分（`scores`），并使用softmax函数将其转换为权重。然后，我们计算`source`和`target`的上下文向量（`source`和`target`），并将它们返回。

值得注意的是，这里的`attention`线性层将输入的`hidden_size*2`维度压缩到了1维，以计算注意力得分。此外，在计算上下文向量时，我们使用了PyTorch中的`bmm`函数来进行批矩阵乘法。

希望这个代码示例能够帮助你更好地理解十字交叉注意力机制的实现方法。