基于Transformer实现机器翻译（日译中）

最新推荐文章于 2024-10-16 01:06:03 发布

JayOoO.

最新推荐文章于 2024-10-16 01:06:03 发布

阅读量1.4k

点赞数 21

文章标签： transformer 机器翻译深度学习

本文链接：https://blog.csdn.net/weixin_72160142/article/details/140061439

版权

Japanese-Chinese Machine Translation Model with Transformer & PyTorch

A tutorial using Jupyter Notebook, PyTorch, Torchtext, and SentencePiece

导入需要的库

import math
import torchtext
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from collections import Counter
from torchtext.vocab import Vocab
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
import io
import time
import pandas as pd
import numpy as np
import pickle
import tqdm
import sentencepiece as spm
torch.manual_seed(0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(torch.cuda.get_device_name(0)) ## 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

读取语料库

# 读取日英平行语料库
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\t', engine='python', header=None)
trainen = df[2].values.tolist()  # 英文句子列表
trainja = df[3].values.tolist()  # 日文句子列表
# trainen.pop(5972)
# trainja.pop(5972)

数据采样

print(trainen[500])
print(trainja[500])

Chinese HS Code Harmonized Code System < HS编码 2905 无环醇及其卤化、磺化、硝化或亚硝化衍生物 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...
Japanese HS Code Harmonized Code System < HSコード 2905 非環式アルコール並びにそのハロゲン化誘導体、スルホン化誘導体、ニトロ化誘導体及びニトロソ化誘導体 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...

我们也可以使用不同的并行数据集来跟随本文，只要确保我们可以将数据处理成如上所示的两个字符串列表，其中包含日语和英语句子。

标记器

与英语或其他按字母顺序排列的语言不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的标记器，它是使用sentencepece为日语和英语创建的

en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='spm.ja.nopretok.model')

加载标记器之后，您可以测试它们，例如，通过执行下面的代码。

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')

['▁All',
 '▁residents',
 '▁aged',
 '▁20',
 '▁to',
 '▁59',
 '▁years',
 '▁who',
 '▁live',
 '▁in',
 '▁Japan',
 '▁must',
 '▁enroll',
 '▁in',
 '▁public',
 '▁pension',
 '▁system',
 '.']

ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type='str')

['▁',
 '年',
 '金',
 '▁日本',
 'に住んでいる',
 '20',
 '歳',
 '~',
 '60',
 '歳の',
 '全ての',
 '人は',
 '、',
 '公的',
 '年',
 '金',
 '制度',
 'に',
 '加入',
 'しなければなりません',
 '。']

建立TorchText词汇对象并将句子转换为Torch张量

def build_vocab(sentences, tokenizer):
    # 构建词汇表
    counter = Counter()
    for sentence in sentences:
        counter.update(tokenizer.encode(sentence, out_type=str))
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
# 构建日语和英语词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)
en_vocab = build_vocab(trainen, en_tokenizer)

在我们有了词汇表对象之后，我们可以使用词汇表和标记器对象来为我们的训练数据构建张量。

def data_process(ja, en):
    # 将句子转换为张量
    data = []
    for (raw_ja, raw_en) in zip(ja, en):
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)], dtype=torch.long)
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)], dtype=torch.long)
        data.append((ja_tensor_, en_tensor_))
    return data
# 处理数据
train_data = data_process(trainja, trainen)

创建要在训练期间迭代的DataLoader对象

在这里，我将BATCH_SIZE设置为16以防止“cuda内存不足”，但这取决于各种事情，例如您的机器内存容量，数据大小等，因此可以根据您的需要随意更改批大小(注意:PyTorch的教程使用Multi30k德语-英语数据集将批大小设置为128)。

BATCH_SIZE = 8
PAD_IDX = ja_vocab['<pad>']
BOS_IDX = ja_vocab['<bos>']
EOS_IDX = ja_vocab['<eos>']
def generate_batch(data_batch):
    # 生成批处理数据
    ja_batch, en_batch = [], []
    for (ja_item, en_item) in data_batch:
        ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
        en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
    ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
    en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)
    return ja_batch, en_batch
# 创建数据加载器
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, collate_fn=generate_batch)

Sequence-to-sequence Transformer

Transformer是在“Attention is all you need”论文中提出的用于解决机器翻译任务的Seq2Seq模型。变压器模型由编码器和解码器块组成，每个块包含固定数量的层。

编码器通过一系列多头注意和前馈网络层对输入序列进行传播处理。编码器的输出称为存储器，与目标张量一起馈送到解码器。编码器和解码器以端到端方式使用teacher force进行培训。

from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)


class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward:int = 512, dropout:float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)

        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))
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,
                                        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)

文本标记通过使用标记嵌入表示。位置编码被添加到标记嵌入中以引入词序的概念。

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        """
        初始化位置编码模块。
        参数：
        emb_size (int): 词嵌入的维度。
        dropout (float): dropout 概率。
        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))
        
        # 计算位置编码：sin 在偶数索引，cos 在奇数索引
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        
        # 增加一个维度以便于与输入的 token embedding 相加
        pos_embedding = pos_embedding.unsqueeze(-2)
        # 定义 dropout 层
        self.dropout = nn.Dropout(dropout)
        
        # 将位置嵌入注册为缓冲区，不作为模型参数
        self.register_buffer('pos_embedding', pos_embedding)
    def forward(self, token_embedding: Tensor):
        """
        前向传播函数。
        参数：
        token_embedding (Tensor): 输入的 token 嵌入。
        返回：
        Tensor: 添加了位置编码后的嵌入。
        """
        # 将位置编码添加到 token 嵌入并应用 dropout
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0), :])
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        """
        初始化 token 嵌入模块。
        参数：
        vocab_size (int): 词汇表大小。
        emb_size (int): 词嵌入的维度。
        """
        super(TokenEmbedding, self).__init__()
        
        # 定义词嵌入层
        self.embedding = nn.Embedding(vocab_size, emb_size)
        
        # 存储词嵌入的维度
        self.emb_size = emb_size
    def forward(self, tokens: Tensor):
        """
        前向传播函数。
        参数：
        tokens (Tensor): 输入的 tokens。
        返回：
        Tensor: 乘以 sqrt(emb_size) 后的词嵌入。
        """
        # 获取 token 的嵌入，并乘以嵌入维度的平方根进行缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续单词掩码来阻止目标单词关注它的后续单词。我们还创建遮罩，用于屏蔽源和目标填充令牌

def generate_square_subsequent_mask(sz):
    """
    生成一个大小为 (sz, sz) 的上三角矩阵掩码，用于掩盖序列中后续的token。
    参数：
    sz (int): 掩码的大小。
    返回：
    mask (Tensor): 一个浮点型的掩码张量，填充为0.0或-inf。
    """
    # 生成一个上三角矩阵，其中上三角部分（包含对角线）为1
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将掩码转换为浮点型，并将0位置填充为负无穷大，1位置填充为0.0
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask
def create_mask(src, tgt):
    """
    创建用于源序列和目标序列的掩码。
    参数：
    src (Tensor): 源序列张量，形状为 (src_seq_len, batch_size)。
    tgt (Tensor): 目标序列张量，形状为 (tgt_seq_len, batch_size)。
    返回：
    src_mask (Tensor): 源序列的掩码张量。
    tgt_mask (Tensor): 目标序列的掩码张量。
    src_padding_mask (Tensor): 源序列的填充掩码张量。
    tgt_padding_mask (Tensor): 目标序列的填充掩码张量。
    """
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]
    # 生成目标序列的掩码
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    
    # 生成源序列的掩码，全为0（因为源序列的掩码不需要考虑序列顺序）
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
    # 生成源序列的填充掩码，src中填充符号的位置为True
    src_padding_mask = (src == PAD_IDX).transpose(0, 1)
    
    # 生成目标序列的填充掩码，tgt中填充符号的位置为True
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
    
    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

# 假设 ja_vocab 和 en_vocab 是分别表示源语言（日语）和目标语言（英语）的词汇表
SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 512  # 嵌入层的维度
NHEAD = 8  # 多头注意力机制的头数
FFN_HID_DIM = 512  # 前馈神经网络的隐藏层维度
BATCH_SIZE = 16  # 批次大小
NUM_ENCODER_LAYERS = 3  # 编码器层数
NUM_DECODER_LAYERS = 3  # 解码器层数
NUM_EPOCHS = 16  # 训练的轮数
# 初始化 Seq2SeqTransformer 模型
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,
                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)
# 初始化模型参数
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)
# 将模型移动到指定设备（GPU或CPU）
transformer = transformer.to(device)
# 定义交叉熵损失函数，忽略填充符号的损失计算
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)
# 定义Adam优化器
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)
def train_epoch(model, train_iter, optimizer):
    """
    训练模型一个epoch。
    参数：
    model (nn.Module): 需要训练的模型。
    train_iter (iterable): 训练数据迭代器。
    optimizer (Optimizer): 优化器。
    返回：
    float: 平均损失。
    """
    model.train()
    losses = 0
    for idx, (src, tgt) in enumerate(train_iter):
        src = src.to(device)
        tgt = tgt.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)
        optimizer.zero_grad()
        # 计算损失，目标序列去掉第一个时间步
        tgt_out = tgt[1:, :]
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
        losses += loss.item()
    return losses / len(train_iter)  
def evaluate(model, val_iter):
    """
    在验证集上评估模型。
    参数：
    model (nn.Module): 需要评估的模型。
    val_iter (iterable): 验证数据迭代器。
    返回：
    float: 平均损失。
    """
    model.eval()
    losses = 0
    with torch.no_grad():  # 评估过程中不计算梯度
        for idx, (src, tgt) in enumerate(val_iter):
            src = src.to(device)
            tgt = tgt.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_iter)

开始训练

最后，在准备好必要的类和函数之后，我们准备训练我们的模型。这是不言而喻的，但是完成训练所需的时间可能会有很大的不同，这取决于很多事情，比如计算能力、参数和数据集的大小。

当我使用JParaCrawl(每种语言大约有590万个句子)的完整句子列表来训练模型时，使用单个NVIDIA GeForce RTX 3070 GPU，每个epoch大约需要5个小时。

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()
  train_loss = train_epoch(transformer, train_iter, optimizer)
  end_time = time.time()
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))

开始训练

首先，我们创建翻译新句子的函数，包括获取日语句子、标记化、转换为张量、推理，然后将结果解码回句子，但这次是英语。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    使用贪婪解码法生成翻译。
    参数：
    model (nn.Module): 序列到序列的Transformer模型。
    src (Tensor): 源序列张量，形状为 (src_seq_len, 1)。
    src_mask (Tensor): 源序列掩码张量。
    max_len (int): 生成序列的最大长度。
    start_symbol (int): 起始符号的索引。
    返回：
    Tensor: 生成的目标序列张量。
    """
    src = src.to(device)
    src_mask = src_mask.to(device)
    
    # 编码源序列
    memory = model.encode(src, src_mask)
    
    # 初始化目标序列，起始符号为 start_symbol
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)
    
    for i in range(max_len - 1):
        memory = memory.to(device)
        
        # 创建 memory 的掩码，全为0，因为 memory 的掩码不需要考虑序列顺序
        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool)
        
        # 生成目标序列的掩码
        tgt_mask = (generate_square_subsequent_mask(ys.size(0))
                    .type(torch.bool)).to(device)
        
        # 解码目标序列
        out = model.decode(ys, memory, tgt_mask)
        
        # 转置输出，使其形状为 (batch_size, seq_len, vocab_size)
        out = out.transpose(0, 1)
        
        # 获取最后一个时间步的输出概率
        prob = model.generator(out[:, -1])
        
        # 使用 torch.max 选择概率最大的词作为下一个词
        _, 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, src, src_vocab, tgt_vocab, src_tokenizer):
    """
    使用模型将源序列翻译成目标语言。
    参数：
    model (nn.Module): 序列到序列的Transformer模型。
    src (str): 源序列字符串。
    src_vocab (Vocab): 源语言的词汇表。
    tgt_vocab (Vocab): 目标语言的词汇表。
    src_tokenizer (Tokenizer): 用于源语言的分词器。
    返回：
    str: 生成的目标语言序列。
    """
    model.eval()
    
    # 将源序列分词并转换为索引序列，添加起始符号和结束符号
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)] + [EOS_IDX]
    num_tokens = len(tokens)
    
    # 将索引序列转换为张量
    src = torch.LongTensor(tokens).reshape(num_tokens, 1)
    
    # 创建源序列的掩码，全为0，因为源序列不需要掩码
    src_mask = torch.zeros(num_tokens, num_tokens).type(torch.bool)
    
    # 使用贪婪解码生成目标序列
    tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
    
    # 将目标序列的索引转换为字符串，并去掉起始符号和结束符号
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

然后，我们可以调用翻译函数并传递所需的参数。

translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)

保存模型

import pickle
# 假设 en_vocab 和 ja_vocab 是词汇表对象
# 保存英语词汇表到文件 'en_vocab.pkl'
with open('en_vocab.pkl', 'wb') as file:
    pickle.dump(en_vocab, file)
# 保存日语词汇表到文件 'ja_vocab.pkl'
with open('ja_vocab.pkl', 'wb') as file:
    pickle.dump(ja_vocab, file)
file.close()

最后，我们还可以使用PyTorch保存和加载函数保存模型以供以后使用。通常，有两种保存模型的方法，这取决于我们以后想要使用它们的目的。第一个仅用于推理，我们可以稍后加载模型并使用它从日语翻译成英语。

# save model for inference
torch.save(transformer.state_dict(), 'inference_model')

第二个也是用于推理的，但当我们稍后想要加载模型并想要恢复训练时也是如此。

# 使用 torch.save() 保存模型和训练检查点
torch.save({
    'epoch': NUM_EPOCHS,                      # 当前训练的轮数
    'model_state_dict': transformer.state_dict(),  # 模型的状态字典，包含所有参数
    'optimizer_state_dict': optimizer.state_dict(),  # 优化器的状态字典，包含优化器的状态信息
    'loss': train_loss,                       # 当前训练的损失值
}, 'model_checkpoint.tar')