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

最新推荐文章于 2024-07-31 01:21:37 发布

听雨悟剑声

最新推荐文章于 2024-07-31 01:21:37 发布

阅读量1.2k

点赞数 44

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

本文链接：https://blog.csdn.net/ai12581/article/details/139975555

版权

一、实验环境

python 3.7

RTX4090D(24GB) AutoDL平台租用

pandas

torchtext==0.6.0

torch==2.1.0（一般>=1.65就没问题，如果使用Auto平台,配置的时候torch直接选2.1.0就行）

sentencepiece==0.1.91

如果显存大于8GB,可以使用RTX3060以上显卡本地运行，作者平台运行差不多一小时三十分钟，一共16轮

二、介绍

在这篇博客中，我将分享一个使用 PyTorch 进行日语到英语翻译模型训练的项目。该项目旨在展示如何从数据处理到模型训练，再到模型保存和加载的完整流程。本项目使用 Transformer 模型，这是近年来在机器翻译任务中表现优异的一种深度学习模型。我们将使用从 JParaCrawl 下载的日英平行数据集。

transformer由Encoder（编码器）和Decoder(解码器）组成，我们将实现transformer的完整流程走一遍。

整个项目分为几个主要步骤：

数据处理: 加载并预处理数据，包括分词和词汇表构建。
模型构建: 构建 Transformer 模型，用于语言翻译。
模型训练: 训练模型，包括设置训练参数和优化策略。
实验结果: 分析训练过程中的关键结果和模型表现。

三、数据处理

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.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 提取第二列（索引为 2）的值并转换为列表
trainen = df[2].values.tolist()

# 提取第三列（索引为 3）的值并转换为列表
trainja = df[3].values.tolist()

# 删除索引为 5972 的元素
# trainen.pop(5972)
# trainja.pop(5972)

我们将数据集中数据集中最后一个有缺失值的数据删除，最终，训练集 trainen 和 trainja 中的句子总数为 5,973,071。

# 使用 SentencePieceProcessor 加载英文模型文件，创建英文 tokenizer
en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')

# 使用 SentencePieceProcessor 加载日文模型文件，创建日文 tokenizer
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)

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

使用SentencePieceProcessor来创建分词器，我们创建了日文分词器和英文分词器，可以输入一句话查看分词效果

结果如下

我们可以看到SentencePieceProcessor分词将空格转义为_,分词效果还可以。

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_vocab 将每个词转换为对应的索引
    ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    # 将英语文本转换为索引张量，并使用 en_vocab 将每个词转换为对应的索引
    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  # 返回处理后的数据列表

# 使用 data_process 函数处理训练数据 trainja 和 trainen，生成训练数据集 train_data
train_data = data_process(trainja, trainen)

对原始训练数据进行处理，生成训练数据集，再将训练数据集使用分词器处理，生成我们需要的词汇表。

BATCH_SIZE = 8  # 定义批处理大小
PAD_IDX = ja_vocab['<pad>']  # 获取词汇表中 <pad> 的索引
BOS_IDX = ja_vocab['<bos>']  # 获取词汇表中 <bos> 的索引
EOS_IDX = ja_vocab['<eos>']  # 获取词汇表中 <eos> 的索引

# 定义批处理生成函数
def generate_batch(data_batch):
  ja_batch, en_batch = [], []  # 初始化日语和英语的批处理列表
  for (ja_item, en_item) in data_batch:  # 遍历数据批次中的每个数据对
    # 为每个日语句子添加 <bos> 和 <eos> 标记，并添加到 ja_batch 列表中
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
    # 为每个英语句子添加 <bos> 和 <eos> 标记，并添加到 en_batch 列表中
    en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
  # 对 ja_batch 列表进行填充，使所有句子长度一致
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
  # 对 en_batch 列表进行填充，使所有句子长度一致
  en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)
  # 返回填充后的日语和英语批处理数据
  return ja_batch, en_batch

# 使用 DataLoader 创建训练数据迭代器
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

我们定义了一个批处理生成函数，并且使用了Dataloder创建了一个训练数据迭代器，多次提供数据，这里如果内存足够可以增大batch_size。

四、模型构建

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

# 定义 Seq2SeqTransformer 类，这是一个序列到序列的 Transformer 模型
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)
        # 定义 Transformer 编码器，由多个编码器层组成
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # 定义解码器层，包含多头注意力机制和前馈神经网络
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 定义 Transformer 解码器，由多个解码器层组成
        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: torch.Tensor, trg: torch.Tensor, src_mask: torch.Tensor,
                tgt_mask: torch.Tensor, src_padding_mask: torch.Tensor,
                tgt_padding_mask: torch.Tensor, memory_key_padding_mask: torch.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: torch.Tensor, src_mask: torch.Tensor):
        # 计算源序列的词嵌入和位置编码，并通过编码器处理
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    # 解码目标序列
    def decode(self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: torch.Tensor):
        # 计算目标序列的词嵌入和位置编码，并通过解码器处理
        return self.transformer_decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)

# 定义 TokenEmbedding 类，用于词嵌入
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size: int):
        super(TokenEmbedding, self).__init__()
        # 定义嵌入层，将词汇表中的词嵌入到指定的维度
        self.embedding = nn.Embedding(vocab_size, emb_size)

    def forward(self, tokens: torch.Tensor):
        # 返回嵌入后的词
        return self.embedding(tokens)

# 定义 PositionalEncoding 类，用于位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout: float):
        super(PositionalEncoding, self).__init__()
        # 定义 Dropout 层
        self.dropout = nn.Dropout(dropout)
        
        # 创建位置编码矩阵
        max_len = 5000
        pe = torch.zeros(max_len, emb_size)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, emb_size, 2).float() * (-torch.log(torch.tensor(10000.0)) / emb_size))
        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: torch.Tensor):
        # 将位置编码添加到输入张量中，并应用 Dropout
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

import torch
import torch.nn as nn
import math
from torch import Tensor

# 定义位置编码类 PositionalEncoding
class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, 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))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 注册位置编码矩阵为模型的缓冲区
        self.register_buffer('pos_embedding', pos_embedding)
        
        # 定义 Dropout 层
        self.dropout = nn.Dropout(dropout)

    # 前向传播函数
    def forward(self, token_embedding: Tensor):
        # 将位置编码添加到词嵌入张量中，并应用 Dropout
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])

# 定义词嵌入类 TokenEmbedding
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):
        # 返回词嵌入张量，并乘以 math.sqrt(self.emb_size)
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

def generate_square_subsequent_mask(sz):
    # 生成一个下三角矩阵，主对角线及其以下的元素为1，其余为0
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 转换为浮点型，并用masked_fill方法进行填充
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return 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)

    # 源序列和目标序列的填充标记掩码
    PAD_IDX = 0  # 假设PAD的索引为0
    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

我们定义了一个序列到序列的 Transformer 模型，采用了Dropout来降低过拟合概论，除此之外我们编写了一个掩码函数用于屏蔽目标序列中的未来位置信息。

五、模型训练

SRC_VOCAB_SIZE = len(ja_vocab)  # 源语言词汇表大小
TGT_VOCAB_SIZE = len(en_vocab)  # 目标语言词汇表大小
EMB_SIZE = 512  # 词嵌入的维度
NHEAD = 8  # 多头注意力机制中头的数量
FFN_HID_DIM = 512  # FeedForward 层的隐藏层维度
BATCH_SIZE = 16  # 批量大小
NUM_ENCODER_LAYERS = 3  # 编码器层数
NUM_DECODER_LAYERS = 3  # 解码器层数
NUM_EPOCHS = 16  # 训练轮数

# 定义 Transformer 模型
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)

# 定义优化器
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)

# 定义训练函数
def train_epoch(model, train_iter, optimizer):
    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
        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.eval()  # 设置模型为评估模式
    losses = 0
    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
        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)  # 返回平均损失

这代码中定义了参数，还定义了一个训练函数和一个评估函数，如果要调参请在这里调试。

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"))

此代码开始训练，我使用RTX4090D训练时，一共16轮，前四轮都在七分钟左右，从第五轮开始每轮训练时间减少，从380s减少到310s，到第十轮开始稳定到310s。

经实测，使用A800(80GB)显卡训练，每轮用时稳定在210s左右。

六、实验结果分析

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(device)  # 将源序列数据移动到指定设备（GPU 或 CPU）
    src_mask = src_mask.to(device)  # 将源序列的掩码数据移动到指定设备
    
    # 编码源序列，得到编码后的记忆（memory）
    memory = model.encode(src, src_mask)
    
    # 初始化目标序列，以起始符号填充（此处起始符号通常为 "<bos>" 的索引）
    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）移动到指定设备
        
        # 创建目标序列的掩码，确保在解码时只能看到已生成部分及其之前的内容
        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)
        
        # 使用 Transformer 模型解码器解码得到输出
        out = model.decode(ys, memory, tgt_mask)
        out = out.transpose(0, 1)  # 转置输出张量的维度
        
        # 通过生成器（generator）生成下一个词的概率分布，并选择概率最高的词作为下一个词
        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, src, src_vocab, tgt_vocab, src_tokenizer):
    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)
    
    # 将源文本转换为张量，并增加一个维度表示批次大小为1
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1))
    
    # 创建源序列的掩码，这里为全零矩阵
    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)

使用翻译函数查看效果

这是网上翻译结果

可以看到翻译效果还是有问题，这是分词器的问题

可以使用另一种分词器构建语料库

sentences = []
with open("datasets/data.txt",encoding="ansi") as file:   #加载文档
    data = file.read().split("\n")
for row in data:        #删除无用字符
    row = row.strip()
    sentences.append(row)
# print(sentence)
with open("datasets/train_data.txt", "w", encoding="utf-8") as file:
    file.write("\n".join(sentences))

因为Sentencepiece中会将空格作为一个特殊的字符，增大了语料的规模，导致训练冗余，所以首先对语料中无效的空格进行删除

import sentencepiece as spm
spm.SentencePieceTrainer.train(
    input='datasets/train_data.txt',
    model_type="bpe",
    model_prefix='tokenizer',   #输出模型名称前缀。训练完成后将生成 <model_name>.model 和 <model_name>.vocab 文件。
    vocab_size=32000,       #目标词表的大小
    user_defined_symbols=['特定称谓', '特定称谓'],#可以设置自己特定的token，并且这个token不会被拆成子词
    character_coverage=1,  #覆盖字符的比例
    max_sentencepiece_length=6  #最大的基本单元,也就是在这个单元内的词进行合并
)

此后再调用其生成分词器即可