【深度】交叉注意力机制

交叉注意力机制，也称为cross-attention，是指在注意力机制中，一个序列中的某个位置与另一个序列中的所有位置进行注意力计算。

import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim):
        super(CrossAttention, self).__init__()
        self.query_dim = query_dim
        self.context_dim = context_dim

        self.linear_q = nn.Linear(query_dim, query_dim)
        self.linear_c = nn.Linear(context_dim, query_dim)

    def forward(self, query, context):
        # Query和Context的维度分别为 [batch_size, query_len, query_dim] 和 [batch_size, context_len, context_dim]
        # 首先将Query和Context分别通过线性变换
        query_proj = self.linear_q(query)  # [batch_size, query_len, query_dim]
        context_proj = self.linear_c(context)  # [batch_size, context_len, query_dim]

        # 计算注意力权重
        attention_weights = torch.bmm(query_proj, context_proj.transpose(1, 2))  # [batch_size, query_len, context_len]
        attention_weights = F.softmax(attention_weights, dim=-1)

        # 对Context序列进行加权求和
        attended_context = torch.bmm(attention_weights, context)  # [batch_size, query_len, context_dim]

        return attended_context, attention_weights