交叉注意力机制(Cross-Attention mechanism)python代码

以下是交叉注意力机制的python代码实现示例：

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.W = nn.Linear(input_dim, input_dim, bias=False)
    
    def forward(self, Q, K, V):
        Q = self.W(Q) # (batch_size, seq_len_q, input_dim)
        K = self.W(K) # (batch_size, seq_len_k, input_dim)
        V = self.W(V) # (batch_size, seq_len_v, input_dim)
        
        # 计算注意力得分
        scores = torch.bmm(Q, K.transpose(1, 2)) # (batch_size, seq_len_q, seq_len_k)
        attn_weights = torch.softmax(scores, dim=-1)
        
        # 加权平均值
        attn_output = torch.bmm(attn_weights, V) # (batch_size, seq_len_q, input_dim)
        
        return attn_output

以上代码实现了一个cross-attention模块，输入Q、K、V均为三维张量，表示query、key、value，分别表示查询序列、键序列、值序列，这三个序列在attention中扮演不同的角色。在函数内部，首先利用线性变换将输入张量的最后一个维度转换为input_dim，然后计算注意力得分，采用softmax函数将得分归一化得到注意力权重，最后将值序列加权平均得到输出。