从代码学习深度学习 - Transformer PyTorch 版

前言

Transformer 模型自 2017 年在论文《Attention is All You Need》中提出以来，彻底改变了自然语言处理（NLP）领域，并在计算机视觉等其他领域展现了强大的潜力。与传统的 RNN 和 LSTM 相比，Transformer 通过自注意力机制（Self-Attention）实现了并行计算，极大地提高了训练效率和模型性能。本博客将通过 PyTorch 实现的 Transformer 模型代码，深入剖析其核心组件，包括多头注意力机制、位置编码、编码器和解码器等。我们将结合代码和文字说明，逐步拆解 Transformer 的实现逻辑，帮助读者从代码层面理解这一经典模型的精髓。

本文基于提供的代码文件（PE.py、EnDecoder.py、MHA.py 和 Transformer.ipynb），完整呈现 Transformer 的 PyTorch 实现，并通过清晰的目录结构和代码注释，带领大家从零开始学习 Transformer 的构建过程。关于训练和可视化部分，这里忽略掉，但是你仍然可以在下面的链接里找到所有的源代码，其中提供了丰富的注释。无论你是深度学习初学者还是希望深入理解 Transformer 的开发者，这篇博客都将为你提供一个清晰的学习路径。

完整代码：下载链接

1. 位置编码（Positional Encoding）

Transformer 的自注意力机制不包含序列的位置信息，因此需要通过位置编码（Positional Encoding）为每个词元添加位置信息。以下是 PE.py 中实现的位置编码类，它通过正弦和余弦函数生成固定位置编码。

import torch
import torch.nn as nn

class PositionalEncoding(nn.Module):
    """位置编码
    
    在Transformer中，由于自注意力机制不含位置信息，需要额外添加位置编码
    在位置嵌入矩阵P中，行代表词元在序列中的位置，列代表位置编码的不同维度
    """
    def __init__(self, num_hiddens, dropout, max_len=1000):
        """初始化位置编码
        
        参数:
            num_hiddens (int): 隐藏层维度，即位置编码的维度
            dropout (float): dropout概率
            max_len (int, 可选): 最大序列长度，默认为1000
        """
        super(PositionalEncoding, self).__init__()
        # 初始化丢弃层
        self.dropout = nn.Dropout(dropout)
        # 创建位置编码矩阵P，形状为(1, max_len, num_hiddens)
        self.P = torch.zeros((1, max_len, num_hiddens))
        # 计算位置编码的正弦和余弦函数输入
        # X形状: (max_len, num_hiddens/2)
        X = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1) / torch.pow(
            10000, torch.arange(0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        # 偶数维度赋值正弦，奇数维度赋值余弦
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)
        
    def forward(self, X):
        """前向传播
        
        参数:
            X (torch.Tensor): 输入张量，形状为(batch_size, seq_len, embed_dim)
            
        返回:
            torch.Tensor: 添加位置编码后的张量，形状为(batch_size, seq_len, embed_dim)
        """
        # 将位置编码加到输入X上，截取与X长度匹配的部分
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        # 应用丢弃并返回结果
        return self.dropout(X)

代码解析：

• 初始化：PositionalEncoding 类根据隐藏层维度（num_hiddens）和最大序列长度（max_len）生成一个位置编码矩阵 P。该矩阵的每一行表示一个位置，每一列对应一个编码维度。
• 正弦和余弦编码：通过正弦（sin）和余弦（cos）函数为不同位置和维度生成编码值，公式为：
  $$PE(pos,2i)=\sin \left(\frac{pos}{10000^{2i/d}}\right),PE(pos,2i+1)=\cos \left(\frac{pos}{10000^{2i/d}}\right)$$
  其中 pos 是位置索引，i 是维度索引，d 是隐藏层维度。
• 前向传播：将输入张量 X 与位置编码矩阵 P 相加，并应用 dropout 以增强模型的鲁棒性。

位置编码的作用是将序列的位置信息嵌入到词嵌入中，使得 Transformer 能够区分相同词元在不同位置的语义。

2. 多头注意力机制（Multi-Head Attention）

多头注意力机制是 Transformer 的核心组件，允许模型并行计算多个注意力头，从而捕获序列中不同方面的依赖关系。以下是 MHA.py 中实现的多头注意力机制。

import math
import torch
from torch import nn
import torch.nn.functional as F

def sequence_mask(X, valid_len, value=0):
    """
    在序列中屏蔽不相关的项，使超出有效长度的位置被设置为指定值
    """
    maxlen = X.size(1)
    mask = torch.arange(maxlen, dtype=torch