attention map 代码

以下是一个简单的自注意力机制（self-attention）的示例代码，可以计算一个序列中每个单词与其他单词的注意力权重（attention weights）：

import torch
import torch.nn.functional as F

class SelfAttention(torch.nn.Module):
    def __init__(self, num_hidden):
        super(SelfAttention, self).__init__()

        self.num_hidden = num_hidden
        self.query = torch.nn.Linear(num_hidden, num_hidden)
        self.key = torch.nn.Linear(num_hidden, num_hidden)
        self.value = torch.nn.Linear(num_hidden, num_hidden)

    def forward(self, inputs):
        # inputs shape: [batch_size, seq_len, num_hidden]

        # Calculate Q, K, V
        q = self.query(inputs)  # [batch_size, seq_len, num_hidden]
        k = self.key(inputs)  # [batch_size, seq_len, num_hidden]
        v = self.value(inputs)  # [batch_size, seq_len, num_hidden]

        # Calculate attention weights
        attn_weights = torch.bmm(q, k.transpose(1, 2))  # [batch_size, seq_len, seq_len]
        attn_weights = F.softmax(attn_weights, dim=-1)

        # Apply attention weights to V
        attn_output = torch.bmm(attn_weights, v)  # [batch_size, seq_len, num_hidden]

        return attn_output, attn_weights

在这个示例代码中，我们定义了一个 `SelfAttention` 类，它接受一个形状为 `[batch_size, seq_len, num_hidden]` 的输入，其中 `batch_size` 表示批次大小，`seq_len` 表示序列长度，`num_hidden` 表示隐藏层的维度。我们使用三个线性层来计算输入序列的查询（query）、键（key）和值（value）向量。然后，我们计算注意力权重（attention weights）并将它们应用于值向量，最后返回注意力输出（attention output）和注意力权重。在计算注意力权重时，我们使用 softmax 函数对每个单词与其他单词的相似度进行归一化处理。