tpa注意力机制python代码

TPA（Triple Perspective Attention）注意力机制的 Python 代码实现如下：

import torch
import torch.nn as nn

class TPAAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(TPAAttention, self).__init__()
        
        self.query_linear = nn.Linear(input_dim, hidden_dim)
        self.key_linear = nn.Linear(input_dim, hidden_dim)
        self.value_linear = nn.Linear(input_dim, hidden_dim)
        
        self.fc = nn.Linear(hidden_dim, 1)
        
    def forward(self, input):
        # input shape: (batch_size, seq_len, input_dim)
        
        # calculate Q,K,V
        Q = self.query_linear(input)  # shape: (batch_size, seq_len, hidden_dim)
        K = self.key_linear(input)  # shape: (batch_size, seq_len, hidden_dim)
        V = self.value_linear(input)  # shape: (batch_size, seq_len, hidden_dim)
        
        # calculate score
        score = torch.tanh(Q + K)  # shape: (batch_size, seq_len, hidden_dim)
        score = self.fc(score).squeeze(dim=2)  # shape: (batch_size, seq_len)
        
        # calculate weight
        weight = torch.softmax(score, dim=1)  # shape: (batch_size, seq_len)
        
        # calculate context vector
        context = torch.bmm(weight.unsqueeze(dim=1), V).squeeze(dim=1)  # shape: (batch_size, hidden_dim)
        
        return context

其中，`input_dim` 为输入特征的维度，`hidden_dim` 为隐藏层的维度。`query_linear`、`key_linear` 和 `value_linear` 分别为线性变换层，用于将输入特征映射到隐藏维度。`fc` 为全连接层，用于计算注意力分数。在 `forward` 方法中，先计算 Q、K、V，然后计算注意力分数 score，并使用 softmax 函数计算得到注意力权重 weight。最后，使用权重 weight 对 V 进行加权求和，得到上下文向量 context。