PyTorch实现多头注意力机制详解：代码示例、注释和变量维度

以下是一个简单的 PyTorch 实现多头注意力机制的代码示例，包括注释和变量维度：

import torch
import torch.nn as nn

# 定义一个多头注意力机制的模块
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        
        assert d_model % num_heads == 0 # 确保能够平均分割
        
        # 定义查询、键、值的线性变换
        self.query_linear = nn.Linear(d_model, d_model)
        self.key_linear = nn.Linear(d_model, d_model)
        self.value_linear = nn.Linear(d_model, d_model)
        
        # 定义输出的线性变换
        self.output_linear = nn.Linear(d_model, d_model)
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 进行线性变换并分割为多头
        query = self.query_linear(query).view(batch_size, -1, self.num_heads, self.d_model // self.num_heads).transpose(1,2) # shape: (batch_size, num_heads, seq_len, d_model // num_heads)
        key = self.key_linear(key).view(batch_size, -1, self.num_heads, self.d_model // self.num_heads).transpose(1,2) # shape: (batch_size, num_heads, seq_len, d_model // num_heads)
        value = self.value_linear(value).view(batch_size, -1, self.num_heads, self.d_model // self.num_heads).transpose(1,2) # shape: (batch_size, num_heads, seq_len, d_model // num_heads)
        
        # 进行注意力计算
        scores = torch.matmul(query, key.transpose(-2, -1)) / (self.d_model // self.num_heads)**0.5 # shape: (batch_size, num_heads, seq_len, seq_len)
        if mask is not None:
            scores = scores.masked_fill(mask==0, -1e9) # 进行mask操作，避免无效信息的影响
        attn_weights = nn.functional.softmax(scores, dim=-1) # shape: (batch_size, num_heads, seq_len, seq_len)
        attn_output = torch.matmul(attn_weights, value) # shape: (batch_size, num_heads, seq_len, d_model // num_heads)
        
        # 将多头的输出拼接起来，并进行线性变换
        concat_output = attn_output.transpose(1,2).contiguous().view(batch_size, -1, self.num_heads * (self.d_model // self.num_heads)) # shape: (batch_size, seq_len, d_model)
        output = self.output_linear(concat_output) # shape: (batch_size, seq_len, d_model)
        
        return output, attn_weights

变量维度的解释：

'query'：查询张量，shape为 '(batch_size, seq_len, d_model)'
'key'：键张量，shape为 '(batch_size, seq_len, d_model)'
'value'：值张量，shape为 '(batch_size, seq_len, d_model)'
'mask'：mask张量，shape为 '(batch_size, seq_len, seq_len)'，其中无效信息的位置为0，有效信息的位置为1
'query_linear'：查询线性变换的全连接层，输入shape为 '(batch_size, seq_len, d_model)'，输出shape为 '(batch_size, seq_len, d_model)'
'key_linear'：键线性变换的全连接层，输入shape为 '(batch_size, seq_len, d_model)'，输出shape为 '(batch_size, seq_len, d_model)'
'value_linear'：值线性变换的全连接层，输入shape为 '(batch_size, seq_len, d_model)'，输出shape为 '(batch_size, seq_len, d_model)'
'query'、'key'、'value'的多头分割结果：shape分别为 '(batch_size, num_heads, seq_len, d_model // num_heads)'
'scores'：注意力分数张量，shape为 '(batch_size, num_heads, seq_len, seq_len)'
'attn_weights'：注意力权重张量，shape为 '(batch_size, num_heads, seq_len, seq_len)'
'attn_output'：多头注意力输出张量，shape为 '(batch_size, num_heads, seq_len, d_model // num_heads)'
'concat_output'：将多头注意力输出拼接后的张量，shape为 '(batch_size, seq_len, d_model)'
'output_linear'：输出线性变换的全连接层，输入shape为 '(batch_size, seq_len, d_model)'，输出shape为 '(batch_size, seq_len, d_model)'
'output'：多头注意力最终输出张量，shape为 '(batch_size, seq_len, d_model)'
'attn_weights'：注意力权重张量，shape为 '(batch_size, num_heads, seq_len, seq_len)'