使用Transformer和PyTorch的中日机器翻译模型

2401_83474403

已于 2024-06-26 09:27:54 修改

阅读量1.3k

点赞数 47

文章标签： transformer pytorch 机器翻译

于 2024-06-23 21:25:17 首次发布

本文链接：https://blog.csdn.net/2401_83474403/article/details/139904959

版权

Transformer模型在自然语言处理领域取得了显著的成功，尤其是在机器翻译任务上。结合PyTorch这一流行的深度学习框架，可以构建高效、准确的日中机器翻译模型。下面将详细分析如何使用Transformer和PyTorch构建日中机器翻译模型：

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')
print(torch.cuda.get_device_name(0)) ## 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

运行上述代码后会显示

device

2.获取平行数据集

数据来源：

可从公共资源如JParaCrawl获取日英平行数据集，尽管需要可能是日中并行数据。
数据处理：使用Pandas读取数据并转换为所需的格式，通常需要将数据分为训练集和测试集。在本教程中，我们将使用从 JParaCrawl 下载的日英并行数据集！[http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl]，它被描述为“NTT创建的最大的公开可用的英日平行语料库。它是通过大量抓取网络并自动对齐平行句子而创建的。

# 读取名为'zh-ja.bicleaner05.txt'的文件，使用制表符('\t')作为分隔符，使用Python引擎解析文件，不包含表头信息
df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 从数据框中提取第3列的数据（索引为2），并将其转换为列表
trainen = df[2].values.tolist()#[:10000]

# 从数据框中提取第4列的数据（索引为3），并将其转换为列表
trainja = df[3].values.tolist()#[:10000]

# 从trainen列表中删除索引为5972的元素
# trainen.pop(5972)

# 从trainja列表中删除索引为5972的元素
# trainja.pop(5972)

在导入所有日语和英语对应数据后，我删除了数据集中的最后一个数据，因为它缺少值。总的来说，trainen 和 trainja 中的句子数为 5,973,071，但是，出于学习目的，通常建议在一次性使用所有数据之前对数据进行采样并确保一切按预期工作，以节省时间。

下面是数据集中包含的句子示例。

# 打印trainen列表的第501个元素（索引为500）
print(trainen[500])
# 打印trainja列表的第501个元素（索引为500）
print(trainja[500])

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

3.准备分词器

分词器的选择：选择适合日语和中文的分词器，通常是基于SentencePiece的分词器。
分词器的应用：使用这些分词器对句子进行编码，以便后续的处理和训练。

与英语或其他字母语言不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的分词器，该分词器是使用SentencePiece创建的日语和英语，您可以访问JParaCrawl网站下载它们，或单击此处。

# 创建一个英文分词器实例，使用预训练的英文模型文件
en_tokenizer = SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')

# 创建一个日文分词器实例，使用预训练的日文模型文件
ja_tokenizer = SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')

加载分词器后，您可以测试它们，例如，通过执行以下代码。

# 对给定的文本进行编码，并将结果转换为字符串类型
encoded_text = 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对给定的文本进行编码，并将结果转换为字符串类型
encoded_text = ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type='str')

4.构建词汇表对象并将句子转换为张量

创建词汇表：通过遍历训练数据和分词结果来构建词汇表。
张量转换：将分词后的句子转换为PyTorch张量，以便于模型的训练和推理。

然后，使用分词器和原始句子，我们构建从 TorchText 导入的 Vocab 对象。此过程可能需要几秒钟或几分钟，具体取决于我们的数据集大小和计算能力。不同的分词器也会影响构建词汇所需的时间，我尝试了其他几种日语分词器，但 SentencePiece 对我来说似乎运行良好且速度足够快。

# 定义一个函数 build_vocab，用于构建词汇表
def build_vocab(sentences, tokenizer):
    # 创建一个计数器对象
    counter = Counter()
    # 遍历输入的句子列表
    for sentence in sentences:
        # 使用分词器对句子进行编码，并更新计数器
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 返回一个词汇表对象，包含特殊符号 '<unk>', '<pad>', '<bos>', '<eos>'
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 使用训练集和日语分词器构建日语词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)
# 使用训练集和英语分词器构建英语词汇表
en_vocab = build_vocab(trainen, en_tokenizer)

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

# 定义一个名为data_process的函数，接收两个参数：ja和en
def data_process(ja, en):
    # 初始化一个空列表data
    data = []
    # 使用zip函数将ja和en中的元素一一对应地组合在一起，然后遍历这些组合
    for (raw_ja, raw_en) in zip(ja, en):
        # 对原始的日语文本进行分词，并将分词结果转换为对应的词汇表索引，然后将这些索引转换为LongTensor类型的张量
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("
"), out_type=str)],
                                dtype=torch.long)
        # 对原始的英语文本进行分词，并将分词结果转换为对应的词汇表索引，然后将这些索引转换为LongTensor类型的张量
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("
"), out_type=str)],
                                dtype=torch.long)
        # 将处理好的日语和英语张量作为元组添加到data列表中
        data.append((ja_tensor_, en_tensor_))
    # 返回处理后的数据列表
    return data

# 调用data_process函数处理训练数据，并将结果赋值给train_data变量
train_data = data_process(trainja, trainen)

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

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

# 设置批量大小为8
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

# 创建训练迭代器，输入为训练数据，批量大小为BATCH_SIZE，打乱顺序，并使用generate_batch函数作为collate_fn参数
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

6.序列到序列转换器

接下来的几个代码和文本说明取自原始的 PyTorch 教程https://pytorch.org/tutorials/beginner/translation_transformer.html。除了BATCH_SIZE之外，我没有做任何更改，de_vocabwhich 这个词被改成了ja_vocab。

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

编码器通过一系列多头注意力和前馈网络层传播输入序列来处理输入序列。编码器的输出称为内存，与目标张量一起馈送到解码器。编码器和解码器使用教师强制技术以端到端的方式进行训练。

# 导入PyTorch中的Transformer相关模块
from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)

# 定义一个Seq2SeqTransformer类，继承自nn.Module
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)
        # 创建Transformer编码器
        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)
        # 创建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: 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)

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

# 定义位置编码类，继承自nn.Module
class PositionalEncoding(nn.Module):
    # 初始化函数，输入参数为emb_size（嵌入维度）、dropout（丢弃率）和maxlen（最大长度，默认为5000）
    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)

        # 初始化dropout层
        self.dropout = nn.Dropout(dropout)
        # 注册位置嵌入矩阵为buffer
        self.register_buffer('pos_embedding', pos_embedding)

    # 前向传播函数，输入参数为token_embedding（词嵌入张量）
    def forward(self, token_embedding: Tensor):
        # 返回经过dropout处理的位置嵌入与词嵌入相加后的结果
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0),:])

# 定义词嵌入类，继承自nn.Module
class TokenEmbedding(nn.Module):
    # 初始化函数，输入参数为vocab_size（词汇表大小）和emb_size（嵌入维度）
    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

    # 前向传播函数，输入参数为tokens（词索引张量）
    def forward(self, tokens: Tensor):
        # 返回经过词嵌入层处理后的词嵌入张量乘以嵌入维度的平方根
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续单词掩码来阻止目标单词关注其后续单词。我们还创建掩码，用于屏蔽源和目标填充tokens。

# 定义一个函数，用于生成平方矩阵的掩码
def generate_square_subsequent_mask(sz):
    # 创建一个上三角矩阵，其中对角线及其以上的元素为1，其余元素为0
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将掩码中的0替换为负无穷大，1替换为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_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
# 设置批量大小
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)

# 定义损失函数，忽略PAD_IDX位置的损失
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 = 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(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)

7.开始训练

在准备了必要的类和函数之后，我们准备训练我们的模型。这是不言而喻的，但完成训练所需的时间可能会有很大差异，具体取决于很多因素，例如计算能力、参数和数据集的大小。

当我使用 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"))

训练配置：设置适当的超参数，如学习率和训练周期数（epoch）。
损失函数与优化器：选择合适的损失函数和优化器来调整模型权重并最小化误差。

8.尝试使用经过训练的模型翻译日语句子

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

# 定义贪婪解码函数
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)
        # 创建目标序列掩码
        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):
    # 对源序列进行编码
    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>", "")

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

# 调用 translate 函数进行翻译
# transformer: 预训练的翻译模型
# text: 要翻译的文本
# ja_vocab: 源语言（日语）的词汇表
# en_vocab: 目标语言（英语）的词汇表
# ja_tokenizer: 用于分割日语文本的工具
translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)

trainen.pop(5)

中文内容

trainja.pop(5)

下面是对应的日语翻译

我又实验了一组中文与日语的翻译

可以看出效果还是相当准确的。

9.保存 Vocab 对象和训练的模型

最后，在训练完成后，我们将首先使用 Pickle 保存 Vocab 对象（en_vocab 和 ja_vocab）。

import pickle
# open a file, where you want to store the data
file = open('en_vocab.pkl', 'wb')
# dump information to that file
pickle.dump(en_vocab, file)
file.close()
file = open('ja_vocab.pkl', 'wb')
pickle.dump(ja_vocab, file)
file.close()

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

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

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

# save model + checkpoint to resume training later
torch.save({
  'epoch': NUM_EPOCHS,
  'model_state_dict': transformer.state_dict(),
  'optimizer_state_dict': optimizer.state_dict(),
  'loss': train_loss,
  }, 'model_checkpoint.tar')

10.代码汇总

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

df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist()#[:10000]
trainja = df[3].values.tolist()#[:10000]
# trainen.pop(5972)
# trainja.pop(5972)
print(trainen[500])
print(trainja[500])
en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')
en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.")
ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。")
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)

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)

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):
        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),:])

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)

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

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

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 = 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(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()
  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):
    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):
    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>", "")

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

trainja.pop(5)

import pickle
# open a file, where you want to store the data
file = open('en_vocab.pkl', 'wb')
# dump information to that file
pickle.dump(en_vocab, file)
file.close()
file = open('ja_vocab.pkl', 'wb')
pickle.dump(ja_vocab, file)
file.close()
# save model for inference
torch.save(transformer.state_dict(), 'inference_model')

# save model + checkpoint to resume training later
torch.save({
  'epoch': NUM_EPOCHS,
  'model_state_dict': transformer.state_dict(),
  'optimizer_state_dict': optimizer.state_dict(),
  'loss': train_loss,
  }, 'model_checkpoint.tar')

11.总结

使用Transformer和PyTorch构建中日机器翻译模型是一项系统的工作，涉及到数据准备、模型构建、训练及测试等多个环节。每一步都需仔细处理以保证最终模型的质量。另外，理解每个步骤的意义及其对最终结果的影响是至关重要的，这不仅有助于构建有效的翻译模型，也为未来可能的优化提供了基础。

2401_83474403

关注

47
点赞
踩
56

收藏

觉得还不错? 一键收藏
0
评论
使用Transformer和PyTorch的中日机器翻译模型

使用Transformer和PyTorch构建日中机器翻译模型是一项系统的工作，涉及到数据准备、模型构建、训练及测试等多个环节。每一步都需仔细处理以保证最终模型的质量。另外，理解每个步骤的意义及其对最终结果的影响是至关重要的，这不仅有助于构建有效的翻译模型，也为未来可能的优化提供了基础。
复制链接

扫一扫