PyTorch 多头注意力机制实现详解及代码示例

PyTorch 实现多头注意力机制

本文将详细介绍如何使用 PyTorch 实现多头注意力机制，并提供完整的代码示例。

代码实现

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

class MultiHeadAttention(nn.Module):
    '多头注意力机制'

    def __init__(self, hidden_size, num_heads):
        super(MultiHeadAttention, self).__init__()

        self.hidden_size = hidden_size
        self.num_heads = num_heads

        self.q_linear = nn.Linear(hidden_size, hidden_size)
        self.k_linear = nn.Linear(hidden_size, hidden_size)
        self.v_linear = nn.Linear(hidden_size, hidden_size)

        self.dropout = nn.Dropout(p=0.1)

        self.fc = nn.Linear(hidden_size, hidden_size)

    def forward(self, q, k, v, mask=None):
        bs = q.size(0)

        # Q、K、V的线性变换
        q = self.q_linear(q).view(bs, -1, self.num_heads, self.hidden_size // self.num_heads).transpose(1, 2)
        k = self.k_linear(k).view(bs, -1, self.num_heads, self.hidden_size // self.num_heads).transpose(1, 2)
        v = self.v_linear(v).view(bs, -1, self.num_heads, self.hidden_size // self.num_heads).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.hidden_size // self.num_heads, dtype=torch.float32))

        # 应用mask
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)

        # 计算注意力权重
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # 计算注意力值
        context = torch.matmul(attention_weights, v)

        # 将多头注意力值拼接在一起
        context = context.transpose(1, 2).contiguous().view(bs, -1, self.hidden_size)

        # 全连接层
        output = self.fc(context)

        return output, attention_weights

模型输入输出

模型输入：
- q: query 向量，形状为 [batch_size, seq_length, hidden_size]
- k: key 向量，形状为 [batch_size, seq_length, hidden_size]
- v: value 向量，形状为 [batch_size, seq_length, hidden_size]
- mask: mask 矩阵，形状为 [batch_size, seq_length]
模型输出：
- output: 注意力输出，形状为 [batch_size, seq_length, hidden_size]
- attention_weights: 注意力权重，形状为 [batch_size, num_heads, seq_length, seq_length]

代码详解

线性变换: 对 query、key 和 value 向量进行线性变换，将它们映射到不同的子空间。
计算注意力分数: 使用 torch.matmul 计算 query 和 key 之间的点积，并除以 sqrt(hidden_size // num_heads) 来进行缩放。
应用mask: 使用 mask 来屏蔽某些位置，例如防止模型关注到填充符号或未来信息。
计算注意力权重: 使用 F.softmax 函数将注意力分数转换为注意力权重。
计算注意力值: 使用注意力权重对 value 向量进行加权求和，得到注意力值。
拼接多头注意力值: 将不同头的注意力值拼接在一起，形成最终的输出。
全连接层: 使用全连接层对输出进行进一步的变换，以获得最终的输出。

总结

本文详细介绍了 PyTorch 中多头注意力机制的实现原理，并提供了完整代码示例。希望对您理解和使用多头注意力机制有所帮助。

如果您有任何问题或建议，请随时在评论区留言！