训练基于Transformer的翻译任务模型

@小森

发布于 2024-12-31 10:37:23

14400

代码可运行

文章被收录于专栏：xiaosen

运行总次数：0

代码可运行

🤗机器翻译是自然语言处理领域中的一个重要任务，其中Transformer模型其强大的表征能力和并行计算能力，成为机器翻译的主流模型。

🤗构建一个基于Transformer的德语到英语翻译模型。

项目使用WMT14德英翻译数据集，该数据集包含约290万对德语-英语句子对。为了加快训练速度，我们选择其中的29000对作为训练集，1000对作为验证集，1000对作为测试集。

模型采用标准的Transformer架构，包括编码器和解码器两部分。编码器和解码器均由多层自注意力机制和前馈神经网络组成。具体参数设置：

编码器层数：3
解码器层数：3
嵌入维度：512
头数：8
前馈维度：512
dropout率：0.1

💫环境配置和硬件检查

import warnings
warnings.filterwarnings("ignore")      
import torch
import torch.nn as nn
import math
from transformers import MarianTokenizer
from datasets import load_dataset
from typing import List
from torch import Tensor
from torch.nn import Transformer
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from timeit import default_timer as timer
import urllib.request
import os
from torch.cuda.amp import GradScaler, autocast
import logging

logging.getLogger("datasets").setLevel(logging.ERROR)   

print("CUDA是否可用:", torch.cuda.is_available())
print("PyTorch版本:", torch.__version__)
if torch.cuda.is_available():
    print("CUDA版本:", torch.version.cuda)

# 设置设备
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("当前使用设备:", DEVICE)
if torch.cuda.is_available():
    print(f"GPU信息: {torch.cuda.get_device_name(0)}")
    print(f"当前GPU显存使用: {torch.cuda.memory_allocated(0)/1024**2:.2f} MB")

💫初始化分词器和定义位置编码

初始化一个用于德语到英语翻译的分词器（Tokenizer）
分词器的作用是将输入的文本转换为模型可以处理的token序列（通常是整数ID），并将模型的输出转换回可读的文本。
为输入的token嵌入添加位置信息。Transformer模型本身没有序列顺序的概念，因此需要通过位置编码来引入序列的位置信息。

tokenizer = MarianTokenizer.from_pretrained('Helsinki-NLP/opus-mt-de-en')

# 定义特殊token的索引
# 填充token的索引，用于将不同长度的序列补齐到相同长度
PAD_IDX = tokenizer.pad_token_id
# 句子开始（Beginning of Sentence）token的索引，用于标记句子的开始
BOS_IDX = tokenizer.bos_token_id
# 句子结束（End of Sentence）token的索引，用于标记句子的结束
EOS_IDX = tokenizer.eos_token_id
# 未知token的索引，用于表示词汇表中不存在的词
UNK_IDX = tokenizer.unk_token_id

# 获取词汇表大小
SRC_VOCAB_SIZE = tokenizer.vocab_size
TGT_VOCAB_SIZE = tokenizer.vocab_size

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout: float, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(-torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        pos_embedding = torch.zeros((maxlen, emb_size))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)
        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])

`💫TokenEmbedding` 类

将输入的 token ID 序列转换为嵌入向量

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

vocab_size：词汇表的大小，即模型需要处理的 token 数量。
emb_size：嵌入向量的维度，即每个 token 被映射为一个多少维的向量。
Token ID 是根据词表（Vocabulary）的大小和内容来分配的。词表是分词器（Tokenizer）中所有可能 token 的集合，每个 token 都会被分配一个唯一的整数 ID

💫构建基于 Transformer 的序列到序列模型

class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, nhead: int, src_vocab_size: int,
                 tgt_vocab_size: int, dim_feedforward: int = 512, dropout: float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        self.transformer = Transformer(d_model=emb_size,
                                    nhead=nhead,
                                    num_encoder_layers=num_encoder_layers,
                                    num_decoder_layers=num_decoder_layers,
                                    dim_feedforward=dim_feedforward,
                                    dropout=dropout)
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)

    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
                tgt_mask: Tensor, src_padding_mask: Tensor,
                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
                              src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
        return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer.encoder(self.positional_encoding(self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer.decoder(self.positional_encoding(self.tgt_tok_emb(tgt)), memory, tgt_mask)

💯这是项目最重要的部分，我详细解释一下：

💢参数：
- num_encoder_layers：编码器的层数。
- num_decoder_layers：解码器的层数。
- emb_size：嵌入向量的维度。
- nhead：多头注意力机制中的头数。
- src_vocab_size：源语言（德语）的词汇表大小。
- tgt_vocab_size：目标语言（英语）的词汇表大小。
- dim_feedforward：前馈神经网络的隐藏层维度。
- dropout：Dropout 率，用于防止过拟合。
💢组件：
- self.transformer：PyTorch 的 Transformer 模块，包含编码器和解码器和前馈神经网络
- self.generator：线性层，将解码器的输出映射为目标语言的词汇表大小，用于生成最终的 token 概率分布。
- self.src_tok_emb：源语言的 token 嵌入层，将源语言的 token ID 转换为嵌入向量。
- self.tgt_tok_emb：目标语言的 token 嵌入层，将目标语言的 token ID 转换为嵌入向量。
- self.positional_encoding：位置编码层，为嵌入向量添加位置信息。

💫生成后续掩码

用于 Transformer 的解码器中，以确保模型在生成目标序列时只能访问当前时刻及之前的信息，而不能访问未来的信息。

def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

💫创建 Transformer 所需的掩码

用于生成 Transformer 模型在训练和推理过程中所需的几种掩码（Mask）

def create_mask(src, tgt):
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]

    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=DEVICE).type(torch.bool)

    src_padding_mask = (src == PAD_IDX).transpose(0, 1)
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

💥 掩码的作用

src_mask：在标准的 Transformer 中，编码器通常不需要掩码（因为编码器可以看到整个输入序列），因此 src_mask 是一个全 False 的矩阵。
tgt_mask: 通过 generate_square_subsequent_mask 生成一个上三角掩码，确保解码器在生成目标序列时只能访问当前时刻及之前的信息，而不能访问未来的信息。
src_padding_mask：填充 token 通常用于将不同长度的序列补齐到相同长度。
tgt_padding_mask: 用于忽略目标序列中的填充 token。

💫数据集下载、加载与处理

def download_multi30k():
    base_url = "https://raw.githubusercontent.com/multi30k/dataset/master/data/task1/raw/"
    
    # 创建数据目录
    os.makedirs("multi30k", exist_ok=True)
    
    # 下载训练、验证和测试数据
    splits = ['train', 'val', 'test']
    languages = ['de', 'en']
    
    for split in splits:
        for lang in languages:
            filename = f"{split}.{lang}"
            url = f"{base_url}{filename}"
            path = f"multi30k/{filename}"
            
            if not os.path.exists(path):
                print(f"Downloading {filename}...")
                urllib.request.urlretrieve(url, path)

def load_data():
    # 加载WMT14数据集的德英对
    dataset = load_dataset("wmt14", "de-en", cache_dir=".cache")
    
    # 为了便于训练，我们只使用一部分数据
    train_size = 29000  # 与Multi30k训练集大小相近
    val_size = 1000
    test_size = 1000
    
    # 处理数据集
    data = {
        'train': {
            'de': [item['de'] for item in dataset['train']['translation'][:train_size]],
            'en': [item['en'] for item in dataset['train']['translation'][:train_size]]
        },
        'val': {
            'de': [item['de'] for item in dataset['validation']['translation'][:val_size]],
            'en': [item['en'] for item in dataset['validation']['translation'][:val_size]]
        },
        'test': {
            'de': [item['de'] for item in dataset['test']['translation'][:test_size]],
            'en': [item['en'] for item in dataset['test']['translation'][:test_size]]
        }
    }
    
    return data

# 添加一个自定义Dataset类
class TranslationDataset(Dataset):
    def __init__(self, de_texts, en_texts):
        self.de_texts = de_texts
        self.en_texts = en_texts
        
    def __len__(self):
        return len(self.de_texts)
    
    def __getitem__(self, idx):
        return {
            'de': self.de_texts[idx],
            'en': self.en_texts[idx]
        }

print("正在加载数据集...")
_cached_data = load_data()  # 全局缓存数据

def get_dataloader(split='train', batch_size=32):
    # 使用缓存的数据而不是重新加载
    data = _cached_data[split]
    
    # 创建Dataset对象
    dataset = TranslationDataset(data['de'], data['en'])
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=(split == 'train')
    )

💫训练模型配置

BATCH_SIZE = 32       
EMB_SIZE = 512        
NHEAD = 8           
FFN_HID_DIM = 512    
NUM_ENCODER_LAYERS = 3  
NUM_DECODER_LAYERS = 3  
NUM_EPOCHS = 18        

# 实例化模型
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
                               NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)
transformer = transformer.to(DEVICE)

# 初始化参数，对模型中所有维度大于1的参数进行Xavier均匀分布初始化。
# 这种初始化方法有助于在训练初期保持梯度的稳定性，避免梯度消失或爆炸。
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

# 定义损失函数和优化器
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
optimizer = torch.optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

# 创建梯度缩放器
scaler = GradScaler()

def train_epoch(model, optimizer):
    try:
        model.train()
        losses = 0
        train_dataloader = get_dataloader('train', BATCH_SIZE)
        
        for batch in train_dataloader:
            src_texts = batch['de']
            tgt_texts = batch['en']
            
            # 使用自动混合精度
            with autocast():
                src_tokens = tokenizer(src_texts, padding=True, return_tensors='pt')
                tgt_tokens = tokenizer(tgt_texts, padding=True, return_tensors='pt')
                
                src = src_tokens['input_ids'].transpose(0, 1).to(DEVICE)
                tgt = tgt_tokens['input_ids'].transpose(0, 1).to(DEVICE)
                
                tgt_input = tgt[:-1, :]
                src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
                
                logits = model(src, tgt_input, src_mask, tgt_mask,
                              src_padding_mask, tgt_padding_mask, src_padding_mask)
                
                tgt_out = tgt[1:, :]
                loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
            
            optimizer.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            losses += loss.item()
        
        return losses / len(train_dataloader)
    except KeyboardInterrupt:
        print("\n训练被手动中断！正在保存当前模型状态...")
        # 保存检查点
        checkpoint = {
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'epoch': epoch,  # 保存当前的epoch
            'train_loss': train_loss,
            'val_loss': val_loss
        }
        torch.save(checkpoint, 'transformer_translation.pth')
        print("模型检查点已保存到 transformer_translation.pth")
        raise KeyboardInterrupt

💫验证集上评估模型的性能

def evaluate(model):
    model.eval()
    losses = 0
    val_dataloader = get_dataloader('val', BATCH_SIZE)
    
    for batch in val_dataloader:
        src_texts = batch['de']
        tgt_texts = batch['en']
        
        src_tokens = tokenizer(src_texts, padding=True, return_tensors='pt')
        tgt_tokens = tokenizer(tgt_texts, padding=True, return_tensors='pt')
        
        src = src_tokens['input_ids'].transpose(0, 1).to(DEVICE)
        tgt = tgt_tokens['input_ids'].transpose(0, 1).to(DEVICE)
        
        tgt_input = tgt[:-1, :]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
        
        logits = model(src, tgt_input, src_mask, tgt_mask,
                      src_padding_mask, tgt_padding_mask, src_padding_mask)
        
        tgt_out = tgt[1:, :]
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        losses += loss.item()
        
    return losses / len(val_dataloader)

💫贪婪解码和翻译功能

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(DEVICE)
    src_mask = src_mask.to(DEVICE)
    
    memory = model.encode(src, src_mask)
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(DEVICE)
    
    for i in range(max_len-1):
        memory = memory.to(DEVICE)
        tgt_mask = (generate_square_subsequent_mask(ys.size(0))
                   .type(torch.bool)).to(DEVICE)
        out = model.decode(ys, memory, tgt_mask)
        out = out.transpose(0, 1)
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.item()
        
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        if next_word == EOS_IDX:
            break
    return ys

def translate(model: torch.nn.Module, src_sentence: str):
    model.eval()
    tokens = tokenizer(src_sentence, return_tensors='pt', padding=True)
    src = tokens['input_ids'].transpose(0, 1).to(DEVICE)
    src_mask = (torch.zeros(src.shape[0], src.shape[0])).type(torch.bool).to(DEVICE)
    
    tgt_tokens = greedy_decode(model, src, src_mask, max_len=src.shape[0] + 5, start_symbol=BOS_IDX).flatten()
    return tokenizer.decode(tgt_tokens.tolist(), skip_special_tokens=True)

greedy_decode 函数用于生成目标语言的翻译结果。它采用贪婪解码策略，即每次选择概率最高的词作为下一个输出，直到生成结束符（EOS_IDX）或达到最大长度。
translate 函数用于将输入的源语言句子翻译成目标语言句子。它调用 greedy_decode 函数生成目标语言序列，并将其解码为可读的文本。

💫模型训练、保存、加载

# 在训练前添加显存清理
if torch.cuda.is_available():
    torch.cuda.empty_cache()



# 训练模型
for epoch in range(1, NUM_EPOCHS + 1):
    start_time = timer()
    train_loss = train_epoch(transformer, optimizer)
    end_time = timer()
    val_loss = evaluate(transformer)
    print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s")

# 保存模型
path = 'transformer_translation.pth'
torch.save(transformer.state_dict(), path)
print("模型保存成功!")

# 加载模型
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
                               NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)
transformer.load_state_dict(torch.load(path))
transformer = transformer.to(DEVICE)
print("模型加载成功!")

# 测试翻译
print(translate(transformer, "Eine Gruppe von Freunden spielt Billiade."))

💦18个Epoch训练了5个小时：