1.导入必要的库
首先，我们要确保我们在python环境中安装了以下库:
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))
Device
用GPU设备来运行这一套代码，设备名称如下：
NVIDIA GeForce GTX 1650
2.平行数据集获取
这里我们将3使用从JParaCrawl下载的日英平行数据集，其下载网址为：http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl，这个数据集被描述为由NTT创建的“最大的公开可用的英日平行语料库”。
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist()#[:10000]# 取出第三列数据作为训练集的英文部分，并转换为列表
trainja = df[3].values.tolist()#[:10000]# 取出第四列数据作为训练集的日文部分，并转换为列表
分词器准备(tokenizers)
与英语或其他字母语言不同，日语句子中不包含用于分隔单词的空格。我们可以使用JParaCrawl提供的分词器(tokenizers)，它使用了SentencePiece分别用于日语和英语。
en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')# 使用 SentencePiece 加载英文模型
ja_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')# 使用SentencePiece 加载日文模型
3.构建 TorchText 的词汇对象并将句子转换为 Torch 张量
利用分词器和原始句子，我们接着构建从 TorchText 导入的词汇（Vocab）对象。这里使用SentencePiece，其经测试工作得足够快且效果良好。
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):
# 使用日文 tokenizer 对当前日文句子进行编码
ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
dtype=torch.long)
# 使用英文 tokenizer 对当前英文句子进行编码
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) #构建训练数据
4.创建 DataLoader 对象以在训练期间进行迭代。
这里根据我的内存容量、数据大小等因素，我将 BATCH_SIZE 设置为 8，以避免 "cuda out of memory" 错误。
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))
# 使用 pad_sequence 函数对日语批次和英语批次进行填充
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
接下来的几段代码和文本解释取自原始的PyTorch教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]。除了BATCH_SIZE和单词de_vocab被更改为ja_vocab之外，没有做任何更改。
Transformer是一种Seq2Seq模型，引入了“Attenion is all you need”论文，用于解决机器翻译任务。Transformer模型由编码器和解码器块组成，每个块包含固定数量的层。
编码器通过一系列的多头注意力和前馈网络层将输入序列进行处理。编码器的输出称为memory，与目标张量一起馈送到解码器中。编码器和解码器使用teacher forcing technique以端到端的方式进行训练。
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))
# 编码器处理源语言序列，生成编码的memoery
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):
# 编码器单独进行编码，返回编码的memorry
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))
# 根据公式计算位置编码中的 sin 与 cos 部分
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):
# 在输入的token_embedding中加入位置编码并进行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):
# 对输入的token进行词嵌入并乘以 sqrt(emb_size) 进行缩放
return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
我们创建一个后续的单词掩码来阻止目标单词关注其后续的单词。我们还创建掩码，用于屏蔽源和目标填充令牌。
def generate_square_subsequent_mask(sz):
# 生成一个上三角矩阵的mask，对角线及以下元素为0，其余为-inf
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]
# 生成目标序列的mask
tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
# 源序列的mask，这里简单地用全零矩阵表示，因为没有用到padding mask
src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
# 生成源序列和目标序列的padding mask
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
定义模型参数并实例化模型。这里当你使用自己的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 # 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) # 将模型移动到指定设备
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)# 定义损失函数，忽略PAD的预测
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, :]# 去掉序列末尾的一个词作为输入
# 创建源序列和目标序列的mask
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, :]
# 创建源序列和目标序列的mask
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)
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):
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函数并传递所需的参数，测试所用的日文如下：“HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)”。
translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)
模型的对上述句子的翻译结果如下：
6.保存词汇对象和训练好的模型。
最后，在培训结束后，我们将首先使用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')
1242
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
