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

最新推荐文章于 2024-06-29 21:25:18 发布

jfbluen

最新推荐文章于 2024-06-29 21:25:18 发布

阅读量392

点赞数 11

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

本文链接：https://blog.csdn.net/jfbluen/article/details/139953767

版权

一、简介

二、具体实现步骤：

1、分词及数据预处理

2、模型构建

3、模型训练

4、模型测试

5、保存Vocab对象和训练好的模型

一、简介：机器翻译（Machine Translation, MT）是指使用计算机自动将一种自然语言（源语言）的文本转换成另一种自然语言（目标语言）的文本的技术。其背后的原理和过程涉及多个关键步骤，包括预处理、编码、解码和训练。下面详细介绍这些步骤及其关联的概念和技术
Transformer是一种用于自然语言处理（NLP）任务的深度学习模型架构，由Vaswani等人在2017年提出。它通过完全依赖于注意力机制而非传统的循环神经网络（RNN）或卷积神经网络（CNN），在多个NLP任务上取得了显著的成功。Transformer模型的核心由编码器（Encoder）和解码器（Decoder）两部分组成，每个部分由多个相同的层堆叠而成。

二、具体实现步骤

1、分词及数据预处理

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

# 读取数据文件，并指定分隔符为制表符
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)

使用JParaCrawl提供的分词器分别对日语和英语文本进行分词处理

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

构建 TorchText Vocab 对象并将句子转换为 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>'])
 
# 示例数据 trainja 和 trainen 分别为日语和英语句子列表
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)

2、模型构建

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:
        # 添加起始标记 <bos> 和结束标记 <eos>，并拼接成张量序列
        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)

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__()
        
        # Transformer 编码器层
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # Transformer 解码器层
        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))
        
        # Transformer 编码器处理源语言序列
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        
        # Transformer 解码器处理目标语言序列
        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):
        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):
        # 返回加上位置编码后的结果，并应用 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):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size
    
    def forward(self, tokens: Tensor):
        # 返回词嵌入乘以 sqrt(词嵌入维度) 的结果
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
 
def generate_square_subsequent_mask(sz):
    """
    生成一个用于 Transformer 解码器的目标序列遮蔽，确保在每个时间步只能看到之前的信息。
    Args:
    - sz (int): 序列的长度
    Returns:
    - mask (Tensor): 形状为 (sz, sz) 的上三角遮蔽矩阵，对角线及以下元素为-∞，其余元素为0
    """
    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
 
def create_mask(src, tgt):
    """
    创建用于 Transformer 模型的遮蔽张量。
    Args:
    - src (Tensor): 源序列张量
    - tgt (Tensor): 目标序列张量
    Returns:
    - src_mask (Tensor): 源序列的填充遮蔽张量，形状为 (src_seq_len, src_seq_len)
    - tgt_mask (Tensor): 目标序列的上三角遮蔽张量，形状为 (tgt_seq_len, tgt_seq_len)
    - src_padding_mask (Tensor): 源序列的填充遮蔽张量，形状为 (src_seq_len, batch_size)
    - tgt_padding_mask (Tensor): 目标序列的填充遮蔽张量，形状为 (tgt_seq_len, batch_size)
    """
    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_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 = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,
                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)
 
# 使用 Xavier 初始化网络参数
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)
 
transformer = transformer.to(device)  # 将模型移动到GPU上（如果可用）
 
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
)

3、模型训练

def train_epoch(model, train_iter, optimizer):
    """
    训练模型的一个epoch。
    Args:
    - model (Seq2SeqTransformer): Transformer 模型
    - train_iter (DataLoader): 训练数据迭代器
    - optimizer (torch.optim.Adam): 优化器
    Returns:
    - 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):
    """
    评估模型在验证集上的表现。
    Args:
    - model (Seq2SeqTransformer): Transformer 模型
    - val_iter (DataLoader): 验证数据迭代器
    Returns:
    - float: 平均验证损失值
    """
    model.eval()
    losses = 0
    for idx, (src, tgt) in (enumerate(valid_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)

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()  # 记录每个epoch的开始时间
  train_loss = train_epoch(transformer, train_iter, optimizer)  # 执行一个epoch的训练并计算损失
  end_time = time.time()  # 记录每个epoch的结束时间
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))  # 打印当前epoch的训练损失和耗时

4、模型测试

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    使用贪婪解码方法生成目标语言序列。
    Args:
    - model (Seq2SeqTransformer): Transformer模型对象。
    - src (Tensor): 源语言输入序列张量，形状为(seq_len, batch_size)。
    - src_mask (Tensor): 源语言输入序列的mask张量，形状为(seq_len, seq_len)。
    - max_len (int): 生成序列的最大长度。
    - start_symbol (int): 目标语言序列的起始符号索引。
    Returns:
    - ys (Tensor): 生成的目标语言序列张量，形状为(seq_len, 1)。
    """
    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)
        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)
        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, src, src_vocab, tgt_vocab, src_tokenizer):
    """
    将源语言文本翻译为目标语言文本。
    Args:
    - model (Seq2SeqTransformer): Transformer模型对象。
    - src (str): 源语言文本字符串。
    - src_vocab (Vocab): 源语言词汇表对象。
    - tgt_vocab (Vocab): 目标语言词汇表对象。
    - src_tokenizer (Tokenizer): 源语言文本分词器。
    Returns:
    - translation (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)
    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>", "")

5、保存Vocab对象和训练好的模型

import pickle
 
# 打开一个文件，用于存储数据
file = open('en_vocab.pkl', 'wb')
 
# 将词汇表对象 en_vocab 序列化并写入文件
pickle.dump(en_vocab, file)
 
# 关闭文件
file.close()
 
# 打开一个文件，用于存储数据
file = open('ja_vocab.pkl', 'wb')
 
# 将词汇表对象 ja_vocab 序列化并写入文件
pickle.dump(ja_vocab, file)
 
# 关闭文件
file.close()
 
# save model for inference
torch.save(transformer.state_dict(), 'inference_model')
 
# 使用torch.save()保存以下内容到'model_checkpoint.tar'文件中：
torch.save({
  'epoch': NUM_EPOCHS,
  'model_state_dict': transformer.state_dict(),
  'optimizer_state_dict': optimizer.state_dict(),
  'loss': train_loss,
  }, 'model_checkpoint.tar')