PyTorch 使用 nn.Transformer 和 TorchText 进行序列到序列建模

在深度学习中，Transformer 架构凭借其卓越的并行性和对序列间长距离依赖关系的精准捕捉，在自然语言处理（NLP）领域大放异彩，从机器翻译到文本生成，Transformer 正在重塑我们对语言处理的认知。对于初学者来说，理解并掌握基于 PyTorch 和 TorchText 的 Transformer 模型构建，是踏入 NLP 高级应用的关键一步。本文将带您从零开始，循序渐进地学习如何利用 PyTorch 的 nn.Transformer 模块和 TorchText 实现序列到序列建模，开启您的 NLP 之旅。

一、环境搭建：筑造模型的基石

在开始构建基于 Transformer 的序列到序列模型之前，确保您的开发环境已正确安装相关的依赖库，这是保证后续代码顺利运行的基础。

1. 安装 PyTorch ：访问Pytorch官方网址，根据您的系统配置（如操作系统、CUDA 版本等）获取适合的安装命令。例如，对于 Linux 系统且 CUDA 版本为 11.8 的用户，可以使用以下命令安装 PyTorch：
   • 命令 ：conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

1. 安装 TorchText ：TorchText 是专门用于处理文本数据的 PyTorch 扩展库，它提供了丰富的数据处理工具和预定义的数据集，极大地方便了文本相关任务的开展。安装命令如下：
   • 命令 ：pip install torchtext

1. 安装 Spacy ：为了实现高效的文本分词操作，本文推荐使用 Spacy 库，它支持多种语言的分词，并且与 TorchText 配合默契。
   • 命令 ：pip install spacy
   • 接下来，您需要下载对应语言的分词模型，例如英语和德语模型：
     • 下载英语模型：python -m spacy download en
     • 下载德语模型：python -m spacy download de

二、数据准备：序列到序列模型的养料

数据是构建模型的基石，对于序列到序列任务，我们选择 Multi30k 数据集，该数据集包含了约 30,000 个英语和德语句子对，句子平均长度约为 13 个单词，非常适合作为训练和评估序列到序列模型的数据基础。

（一）数据集加载与字段定义

TorchText 提供了便捷的数据集加载方式，能够轻松加载 Multi30k 数据集并进行初步处理。同时，我们需要定义字段，指定对源语言（德语）和目标语言（英语）的预处理方式，如分词、添加起始和结束标记等。

from torchtext.datasets import Multi30k
from torchtext.data import Field
## 定义德语字段（源语言）
SRC = Field(tokenize="spacy", tokenizer_language="de", init_token="<sos>", eos_token="<eos>", lower=True)
## 定义英语字段（目标语言）
TRG = Field(tokenize="spacy", tokenizer_language="en", init_token="<sos>", eos_token="<eos>", lower=True)
## 加载 Multi30k 数据集
train_data, valid_data, test_data = Multi30k(language_pair=("de", "en"))

（二）词汇表构建与数据迭代器创建

在对数据进行编码和解码之前，我们需要构建词汇表，将文本单词映射到数值索引。此外，为了高效地将数据喂入模型进行训练和评估，我们还需要使用数据迭代器。TorchText 提供的 BucketIterator 能够根据序列长度将相似长度的样本划分到一个批次中，从而减少填充操作，提高训练效率。

from torchtext.data import BucketIterator
## 构建德语词汇表
SRC.build_vocab(train_data, min_freq=2)
## 构建英语词汇表
TRG.build_vocab(train_data, min_freq=2)
## 设置设备（GPU 或 CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
## 定义批次大小
BATCH_SIZE = 128
## 创建训练、验证和测试数据迭代器
train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data),
    batch_size=BATCH_SIZE,
    device=device
)

三、模型构建：序列到序列的核心引擎

（一）Transformer 模型架构概览

Transformer 模型主要由编码器（Encoder）和解码器（Decoder）组成。编码器负责将输入序列转换为上下文表示，解码器则基于编码器的输出生成目标序列。Transformer 的核心创新在于自注意力机制（Self-Attention），它允许模型在处理序列中的每个位置时，动态地关注序列中其他位置的相关信息，从而捕捉序列中的长距离依赖关系。

（二）编码器与解码器的实现

PyTorch 提供了 nn.TransformerEncoder 和 nn.TransformerDecoder 模块，简化了 Transformer 编码器和解码器的构建过程。以下是编码器和解码器的代码实现：

import torch.nn as nn
class TransformerSeq2Seq(nn.Module):
    def __init__(self, src_vocab_size, trg_vocab_size, emb_dim, hid_dim, n_layers, n_heads, dropout, device, max_len=100):
        super().__init__()
        self.device = device
        
        # 编码器嵌入层和位置编码
        self.src_embedding = nn.Embedding(src_vocab_size, emb_dim)
        self.src_pos_encoder = PositionalEncoding(emb_dim, dropout, max_len)
        
        # 编码器层
        encoder_layer = nn.TransformerEncoderLayer(d_model=emb_dim, nhead=n_heads, dim_feedforward=hid_dim, dropout=dropout)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        
        # 解码器嵌入层和位置编码
        self.trg_embedding = nn.Embedding(trg_vocab_size, emb_dim)
        self.trg_pos_encoder = PositionalEncoding(emb_dim, dropout, max_len)
        
        # 解码器层
        decoder_layer = nn.TransformerDecoderLayer(d_model=emb_dim, nhead=n_heads, dim_feedforward=hid_dim, dropout=dropout)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_layers)
        
        # 输出层
        self.out = nn.Linear(emb_dim, trg_vocab_size)
        
    def forward(self, src, trg, src_mask=None, trg_mask=None, src_padding_mask=None, trg_padding_mask=None):
        # 编码器部分
        src_emb = self.src_embedding(src) * math.sqrt(self.emb_dim)
        src_emb = self.src_pos_encoder(src_emb)
        memory = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
        
        # 解码器部分
        trg_emb = self.trg_embedding(trg) * math.sqrt(self.emb_dim)
        trg_emb = self.trg_pos_encoder(trg_emb)
        output = self.decoder(trg_emb, memory, tgt_mask=trg_mask, memory_key_padding_mask=src_padding_mask, tgt_key_padding_mask=trg_padding_mask)
        
        # 输出层
        output = self.out(output)
        return output

（三）位置编码模块

位置编码模块用于向模型注入序列中单词的位置信息，因为 Transformer 模型本身并不像 RNN 那样对序列顺序敏感。位置编码的维度与嵌入维度相同，因此可以将两者相加。在这里，我们使用不同频率的 sine 和 cosine 函数来实现位置编码。

import math
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer("pe", pe)
        
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

四、模型训练与评估：打磨序列到序列模型的利器

（一）超参数设置与模型初始化

在训练模型之前，我们需要设置一些超参数，这些超参数将影响模型的训练过程和最终性能。接下来，我们根据超参数初始化编码器、解码器和序列到序列模型，并定义优化器来更新模型参数。

## 超参数设置
SRC_VOCAB_SIZE = len(SRC.vocab)
TRG_VOCAB_SIZE = len(TRG.vocab)
EMB_DIM = 512
HID_DIM = 1024
N_LAYERS = 3
N_HEADS = 8
DROPOUT = 0.1
## 初始化模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TransformerSeq2Seq(SRC_VOCAB_SIZE, TRG_VOCAB_SIZE, EMB_DIM, HID_DIM, N_LAYERS, N_HEADS, DROPOUT, device).to(device)
## 定义优化器
optimizer = torch.optim.Adam(model.parameters())

（二）损失函数定义与训练过程

为了评估模型的性能并指导模型训练，我们需要定义损失函数。在序列到序列任务中，通常使用交叉熵损失函数，并忽略填充部分的损失计算。训练过程是模型学习数据模式并不断提升翻译性能的关键阶段。在每个训练周期（epoch）中，模型会处理整个训练数据集，并根据计算得到的损失更新模型参数。

import torch
import math
import time
## 定义损失函数
PAD_IDX = TRG.vocab.stoi["<pad>"]
criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0
    for _, batch in enumerate(iterator):
        src = batch.src
        trg = batch.trg
        
        optimizer.zero_grad()
        
        # 创建掩码
        src_mask = (src != SRC.vocab.stoi["<pad>"]).unsqueeze(-2)
        trg_mask = (trg != TRG.vocab.stoi["<pad>"]).unsqueeze(-2)
        trg_mask = trg_mask & (torch.triu(torch.ones(trg.shape[0], trg.shape[0])) == 0).to(trg_mask.dtype).to(trg_mask.device)
        
        output = model(src, trg[:-1], src_mask=src_mask, trg_mask=trg_mask)
        
        output = output.contiguous().view(-1, output.shape[-1])
        trg = trg[1:].contiguous().view(-1)
        
        loss = criterion(output, trg)
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        
        epoch_loss += loss.item()
    return epoch_loss / len(iterator)
def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0
    with torch.no_grad():
        for _, batch in enumerate(iterator):
            src = batch.src
            trg = batch.trg
            
            src_mask = (src != SRC.vocab.stoi["<pad>"]).unsqueeze(-2)
            trg_mask = (trg != TRG.vocab.stoi["<pad>"]).unsqueeze(-2)
            trg_mask = trg_mask & (torch.triu(torch.ones(trg.shape[0], trg.shape[0])) == 0).to(trg_mask.dtype).to(trg_mask.device)
            
            output = model(src, trg[:-1], src_mask=src_mask, trg_mask=trg_mask)
            
            output = output.contiguous().view(-1, output.shape[-1])
            trg = trg[1:].contiguous().view(-1)
            
            loss = criterion(output, trg)
            
            epoch_loss += loss.item()
    return epoch_loss / len(iterator)
## 训练模型
N_EPOCHS = 10
CLIP = 1
best_valid_loss = float("inf")
for epoch in range(N_EPOCHS):
    start_time = time.time()
    
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, valid_iterator, criterion)
    
    end_time = time.time()
    
    epoch_mins, epoch_secs = divmod(end_time - start_time, 60)
    
    print(f"Epoch: {epoch + 1:02} | Time: {epoch_mins}m {epoch_secs:.2f}s")
    print(f"\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}")
    print(f"\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}")
    
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), "transformer_model.pt")
## 加载最佳模型并在测试集上评估
model.load_state_dict(torch.load("transformer_model.pt"))
test_loss = evaluate(model, test_iterator, criterion)
print(f"| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |")

五、总结与展望

通过本文，您已经学习了如何利用 PyTorch 和 TorchText 构建一个基于 Transformer 的序列到序列模型，从环境搭建、数据准备、模型构建到训练与评估，每一步都至关重要。Transformer 模型凭借其强大的并行计算能力和对长距离依赖关系的捕捉能力，在众多序列到序列任务中表现出色。在实际应用中，您可以根据需求进一步优化模型，如调整超参数、使用更大的模型架构、采用数据增强技术等，以提升模型性能。

未来，随着深度学习技术的不断发展，Transformer 架构及其变体将在自然语言处理领域发挥更加重要的作用，为人们的语言交流和信息处理提供更强大的支持。编程狮将持续为您提供更多优质的技术教程和资源，助力您的编程学习之旅。