机器翻译实例：基于Transformer和PyTorch的中日机器翻译模型

最新推荐文章于 2024-08-06 20:19:36 发布

Goldenworker

最新推荐文章于 2024-08-06 20:19:36 发布

阅读量770

点赞数 28

文章标签： transformer pytorch 机器翻译

本文链接：https://blog.csdn.net/FanShelby/article/details/139995358

版权

本文将简单介绍，如何基于transformer和pytorch，实现一个中日机器翻译模型。

1.导入需要的库

首先，我们定义一个基于Transformer和PyTorch的中日机器翻译模型：

（1）导入了一些必要的库，然后定义模型的参数，如词汇表大小、嵌入维度、编码器层数等。

（2）定义了一个Encoder类，用于实现Transformer的编码器部分，包括多头自注意力机制和前馈神经网络。

（3）定义了一个Decoder类，用于实现Transformer的解码器部分，包括多头自注意力机制、编码器-解码器注意力机制和前馈神经网络。

（4）定义了一个Seq2Seq类，将编码器和解码器组合在一起，实现了一个序列到序列的模型。

代码如下：

# 导入所需库
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)

# 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 如果你有GPU，请在你自己的电脑上尝试运行这一套代码
# print(torch.cuda.get_device_name(0))

2.获取平行数据集

在本实验中，我们将使用从JParaCrawl下载的日英平行数据集（http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl），该数据集被描述为“由NTT创建的最大的公开可用的英日平行语料库。它主要是通过大量爬取网络并自动对齐平行句子来创建的。你也可以在这里查看论文。

代码如下：

# 导入pandas库
import pandas as pd

# 读取数据集文件，使用制表符分隔，设置engine为'python'，不包含表头
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 将第3列的数据（英文数据）转换为列表，并赋值给trainen
trainen = df[2].values.tolist()

# 将第4列的数据（日文数据）转换为列表，并赋值给trainja
trainja = df[3].values.tolist()

# 注释掉的代码表示删除索引为5972的元素，但这里已经注释掉了，所以不会执行
# trainen.pop(5972)
# trainja.pop(5972)

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

以下是数据集中包含的一个句子示例。

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

运行结果如下：

3.准备分词器

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

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.", out_type='str')

运行结果：

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

运行结果：

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

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

代码如下：

# 导入所需库
from collections import Counter
from torchtext.vocab import Vocab
import 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("\
"), out_type=str)],
                                 dtype=torch.long)  # 将索引列表转换为张量

        # 将英语文本转换为标记列表，去除末尾的换行符，并获取对应的词汇表索引
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\
"), out_type=str)],
                                 dtype=torch.long)  # 将索引列表转换为张量

        data.append((ja_tensor_, en_tensor_))  # 将日语和英语的张量对添加到数据列表中
    return data  # 返回处理后的数据列表

# 使用处理函数处理训练数据
train_data = data_process(trainja, trainen)  # 处理训练数据，得到训练数据集的张量表示

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

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

代码如下：

# 设置批处理大小为16，以防止“cuda内存不足”错误
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

# 创建DataLoader对象，用于训练迭代
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

6.Sequence-to-sequence Transformer

6.1Transformer的序列到序列（Seq2Seq）模型

我们定义一个基于Transformer的序列到序列（Seq2Seq）模型，用于机器翻译等任务。Seq2Seq模型包括编码器和解码器两个部分，分别对源语言和目标语言进行编码和解码。

具体步骤为：

首先，从torch.nn模块中导入了TransformerEncoder、TransformerDecoder、TransformerEncoderLayer和TransformerDecoderLayer类。
定义了一个名为Seq2SeqTransformer的类，继承自nn.Module。这个类的构造函数接收以下参数：
- num_encoder_layers: 编码器层数
- num_decoder_layers: 解码器层数
- emb_size: 词嵌入向量的维度
- src_vocab_size: 源语言词汇表大小
- tgt_vocab_size: 目标语言词汇表大小
- dim_feedforward: 前馈神经网络的隐藏层维度，默认为512
- dropout: dropout概率，默认为0.1
在Seq2SeqTransformer类的构造函数中，首先创建了编码器层和解码器层，然后创建了编码器和解码器。接着创建了一个生成器（线性层），用于将编码器的输出转换为目标词汇表的大小。最后，创建了源语言和目标语言的词嵌入层以及位置编码层。
Seq2SeqTransformer类还定义了三个方法：forward、encode和decode。
- forward方法接收源语言输入、目标语言输入以及各种掩码（如src_mask、tgt_mask等），首先对输入进行词嵌入和位置编码，然后使用编码器处理源语言输入，接着使用解码器处理目标语言输入和编码器的输出，最后使用生成器将解码器的输出转换为目标词汇表的大小。
- encode方法接收源语言输入和掩码，对输入进行词嵌入和位置编码，然后使用编码器进行处理。
- decode方法接收目标语言输入、编码器的输出和掩码，对输入进行词嵌入和位置编码，然后使用解码器进行处理。

代码如下：

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)

6.2 位置编码类（PositionalEncoding）和词嵌入类（TokenEmbedding）

接着，我们定义了两个类，分别是位置编码类（PositionalEncoding）和词嵌入类（TokenEmbedding）。

位置编码类的作用：为输入的词嵌入向量添加位置信息。

它首先计算位置编码的分母，然后生成位置序列，接着初始化位置嵌入矩阵，并计算位置嵌入矩阵的偶数列和奇数列。最后，将位置嵌入矩阵与词嵌入相加，并进行dropout操作。

词嵌入类的作用：将输入的tokens转换为词嵌入向量。

它首先初始化一个词嵌入层，然后在forward方法中对输入的tokens进行词嵌入，并乘以emb_size的平方根进行缩放。

代码如下：

# 定义位置编码类
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):
        # 对输入的tokens进行词嵌入，并乘以emb_size的平方根进行缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

6.3单词掩码

为了阻止目标词关注其后的单词，我们创建了一个后续单词掩码。同时，我们还创建了掩码，用于掩盖源和目标填充令牌。

我们代码定义了两个函数，用于生成掩码。

函数一：generate_square_subsequent_mask(sz)

用于生成下三角矩阵，用于屏蔽目标词后面的单词。它使用torch.triu生成一个上三角矩阵，然后转置得到下三角矩阵。接着将矩阵中的0和1分别替换为负无穷和0。

函数二：create_mask(src, tgt)

用于生成源语言和目标语言的掩码以及填充掩码。首先获取源语言和目标语言的长度，然后调用generate_square_subsequent_mask函数生成目标语言的掩码。接着创建一个全零矩阵作为源语言的掩码。最后创建源语言和目标语言的填充掩码，并返回源语言和目标语言的掩码以及填充掩码。

代码如下：

# 创建一个生成下三角矩阵的函数，用于屏蔽目标词后面的单词
def generate_square_subsequent_mask(sz):
    # 使用torch.triu生成一个上三角矩阵，然后转置得到下三角矩阵
    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

6.4定义模型参数并实例化模型

注：由于服务器计算能力有限，想要看到训练的效果请使用自己的带GPU的电脑运行这一套代码。当你使用自己的GPU的时候，NUM_ENCODER_LAYERS 和 NUM_DECODER_LAYERS 设置为3或者更高，NHEAD设置8，EMB_SIZE设置为512

代码如下：

# 定义模型参数并实例化模型
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)

# 将模型放到设备上（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
)

# 定义训练一个epoch的函数
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(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)  # 返回平均损失

7.开始训练

最后，在准备好必要的类和函数之后，我们准备训练我们的模型。不用说，完成训练所需的时间可能会因许多因素而有很大差异，例如计算能力、参数和数据集的大小。

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

代码如下:

# 开始训练
# 最后，在准备好必要的类和函数后，我们准备训练我们的模型。不用说，完成训练所需的时间可能会有很大差异，这取决于许多因素，如计算能力、参数和数据集的大小。

# 当我使用JParaCrawl的完整句子列表（每种语言约590万个句子）进行训练时，每个epoch大约需要5小时，使用单个NVIDIA GeForce RTX 3070 GPU。

# 以下是代码：

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()  # 记录当前时间
  train_loss = train_epoch(transformer, train_iter, optimizer)  # 调用train_epoch函数进行训练，并获取训练损失
  end_time = time.time()  # 记录当前时间
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))  # 打印每个epoch的训练损失和耗时

注：由于服务器计算能力有限，想要看到训练的效果请使用自己的带GPU的电脑运行这一套代码。

8.尝试使用训练好的模型翻译一个日语句子

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

代码如下:

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(device)  # 将源句子转换为设备（GPU或CPU）上的张量
    src_mask = src_mask.to(device)  # 将源句子掩码转换为设备（GPU或CPU）上的张量
    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)  # 将编码后的源句子转换为设备（GPU或CPU）上的张量
        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.
    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)

9.保存词汇对象和训练好的模型

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

import pickle
# 导入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()

最后，我们可以使用PyTorch的保存和加载功能将模型保存以供以后使用。通常，根据我们稍后想要使用的目的，有两种方法可以保存模型。第一种是仅用于推理，我们可以稍后加载模型并使用它从日语翻译成英语。

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

第二种方法也是用于推理，但也适用于稍后加载模型并希望继续训练的情况。

# 保存模型和检查点以便于后续恢复训练
torch.save({
  'epoch': NUM_EPOCHS,
  'model_state_dict': transformer.state_dict(),
  'optimizer_state_dict': optimizer.state_dict(),
  'loss': train_loss,
  }, 'model_checkpoint.tar')
# 保存模型、优化器状态字典、损失值和训练轮数等信息，以便在之后恢复训练

至此结束，再想看也没了。

Goldenworker

关注

28
点赞
踩
11

收藏

觉得还不错? 一键收藏
0
评论
机器翻译实例：基于Transformer和PyTorch的中日机器翻译模型

注：由于服务器计算能力有限，想要看到训练的效果请使用自己的带GPU的电脑运行这一套代码。当你使用自己的GPU的时候，NUM_ENCODER_LAYERS 和 NUM_DECODER_LAYERS 设置为3或者更高，NHEAD设置8，EMB_SIZE设置为512。
复制链接

扫一扫