目录
一、基本概述
基于注意力机制的机器翻译(Attention-based Machine Translation)是神经机器翻译(Neural Machine Translation, NMT)的一个重要发展方向。注意力机制允许翻译模型在解码过程中动态地关注源语言序列的不同部分,从而显著提高翻译质量和模型效率。
1. 背景
传统的编码器-解码器(Encoder-Decoder)模型在处理长句子时会遇到困难,因为编码器需要将整个源语言序列压缩成一个固定长度的上下文向量,这可能导致信息丢失。注意力机制通过引入动态权重,使模型能够在生成每个目标词时关注源语言序列的不同部分,从而缓解了这一问题。
注意力机制存在一些优势:① 动态权重:注意力机制允许模型在每个解码步骤中动态调整权重,更好地捕捉上下文信息。② 处理长序列:缓解了传统编码器-解码器模型在处理长句子时的信息压缩问题。③ 提高翻译质量:通过关注源语言序列的不同部分,生成更自然和准确的翻译。
2. 注意力机制的基本概念
注意力机制的核心思想是为每个目标词生成时,计算一组权重,这些权重表示源语言序列中各个位置的重要性。具体来说,注意力机制通过以下步骤实现:
1.计算注意力权重:对每个源语言词和当前解码状态,计算一个注意力得分。
2.归一化权重:使用softmax函数将注意力得分归一化为概率分布。
3.加权求和:将源语言词的隐状态向量按照注意力权重加权求和,得到上下文向量。
3. 注意力机制的实现
3.1. 全局注意力(Global Attention)
全局注意力机制在每个时间步都考虑整个源语言序列。常用的全局注意力机制有Bahdanau Attention和Luong Attention:
(1) Bahdanau Attention:也称为加性注意力(Additive Attention),计算注意力得分时使用一个前馈神经网络。其中,S(i−1) 是解码器的前一时刻的隐状态,hj 是编码器的隐状态,Wa 和 va 是可训练的权重矩阵和向量。
(2) Luong Attention:也称为乘性注意力(Multiplicative Attention),计算注意力得分时使用点积操作。其中, Wa 是可训练的权重矩阵。
或
3.2. 局部注意力(Local Attention)
局部注意力的关键在于确定每个解码步骤中需要关注的源语言序列的子集。一般来说,这个子集是以解码器当前状态为中心的一段固定长度的窗口。
· 确定窗口中心位置:首先根据解码器当前状态 s{i-1} 估计窗口的中心位置。一个常用的方法是通过位置预测函数来预测中心位置 pi:pi=round(i⋅S/T)。其中, i 是当前解码步骤, T是目标序列的总长度, S是源序列的总长度。
· 定义窗口范围:以pi为中心,定义一个长度为2D的窗口范围 [pi−D,pi+D],其中D是窗口的半宽。
· 计算局部注意力权重:仅计算窗口范围内的源语言序列位置的注意力权重。对于窗口外的源语言位置,注意力权重设置为0。使用softmax函数对窗口内的注意力得分进行归一化。
· 生成上下文向量:根据归一化后的注意力权重,计算窗口范围内源语言隐状态的加权和,生成上下文向量。
4. Transformer模型
Transformer模型是目前最先进的基于注意力机制的机器翻译模型,完全摆脱了RNN的结构。其由编码器和解码器两部分组成,每部分又由多个相同的层堆叠而成。典型的Transformer模型包含6个编码器层和6个解码器层。Transformer模型不仅在机器翻译中表现出色,还广泛应用于其他NLP任务,如文本生成、文本摘要、问答系统等。此外,基于Transformer的变体,如BERT、GPT等,在各类NLP任务中都取得了显著的进展。
二、编码器-解码器
1. 编码器(Encoder)
编码器的作用是把一个不定长的输入序列变换成一个定长的背景变量𝑐,并在该背景变量中编码输入序列信息。常用的编码器是循环神经网络。
考虑批量大小为1的时序数据样本。假设输入序列是𝑥1,…,𝑥𝑇,例如𝑥𝑖是输入句子中的第𝑖个词。在时间步𝑡,循环神经网络将输入𝑥𝑡的特征向量𝑥𝑡和上个时间步的隐藏状态ℎ𝑡−1变换为当前时间步的隐藏状态ℎ𝑡。可用函数𝑓表达循环神经网络隐藏层的变换:ℎ𝑡=𝑓(𝑥𝑡,ℎ𝑡−1).
接下来,编码器通过自定义函数𝑞将各个时间步的隐藏状态变换为背景变量𝑐=𝑞(ℎ1,…,ℎ𝑇).当选择𝑞(ℎ1,…,ℎ𝑇)=ℎ𝑇时,背景变量是输入序列最终时间步的隐藏状态ℎ𝑇。
以上描述的编码器是一个单向的循环神经网络,每个时间步的隐藏状态只取决于该时间步及之前的输入子序列。我们也可以使用双向循环神经网络构造编码器。在这种情况下,编码器每个时间步的隐藏状态同时取决于该时间步之前和之后的子序列(包括当前时间步的输入),并编码了整个序列的信息。
2. 解码器(Decoder)
给定训练样本中的输出序列𝑦1,𝑦2,…,𝑦𝑇′,对每个时间步𝑡′,解码器输出𝑦𝑡′的条件概率将基于之前的输出序列𝑦1,…,𝑦𝑡′−1和背景变量𝑐,即𝑃(𝑦𝑡′∣𝑦1,…,𝑦𝑡′−1,𝑐)。可使用另一个循环神经网络作为解码器。在输出序列的时间步𝑡′,解码器将上一时间步的输出𝑦𝑡′−1以及背景变量𝑐作为输入,并将它们与上一时间步的隐藏状态𝑠𝑡′−1变换为当前时间步的隐藏状态𝑠𝑡′𝑠𝑡′。因此,可函数𝑔𝑔表达解码器隐藏层的变换:𝑠𝑡′=𝑔(𝑦𝑡′−1,𝑐,𝑠𝑡′−1).𝑠𝑡′=𝑔(𝑦𝑡′−1,𝑐,𝑠𝑡′−1).
有了解码器的隐藏状态后,可使用自定义的输出层和softmax运算来计算 𝑃(𝑦𝑡′∣𝑦1,…,𝑦𝑡′−1,𝑐),基于当前时间步的解码器隐藏状态 𝑠𝑡′、上一时间步的输出𝑦𝑡′−1以及背景变量𝑐来计算当前时间步输出𝑦𝑡′的概率分布。
三、编码器—解码器预测不定长序列
1. 贪婪搜索
对于输出序列任一时间步𝑡′,我们从|Y|个词中搜索出条件概率最大的词 𝑦𝑡′ = argmax (𝑦∈𝑌) 𝑃(𝑦|𝑦1,…,𝑦𝑡′−1,𝑐)作为输出。一旦搜索出"<eos>"符号,或者输出序列长度已经达到了最大长度𝑇′𝑇′,便完成输出。
假设输出词典里面有“A”“B”“C”和“<eos>”这4个词。每个时间步下的4个数字分别代表了该时间步生成“A”“B”“C”和“<eos>”这4个词的条件概率。在每个时间步,贪婪搜索选取条件概率最大的词。因此,图中将生成输出序列“A”“B”“C”“<eos>”。该输出序列的条件概率为0.5×0.4×0.4×0.6=0.0480.5×0.4×0.4×0.6=0.048。
接下来,在时间步2中选取了条件概率第二大的词“C”。由于时间步3所基于的时间步1和2的输出子序列由上图中的“A”“B”变为了下图中的“A”“C”,下图中时间步3生成各个词的条件概率发生了变化。我们选取条件概率最大的词“B”。此时时间步4所基于的前3个时间步的输出子序列为“A”“C”“B”,与上图中的“A”“B”“C”不同。因此,下图中时间步4生成各个词的条件概率也与上图中的不同。此时的输出序列“A”“C”“B”“<eos>”的条件概率是0.5×0.3×0.6×0.6=0.0540.5×0.3×0.6×0.6=0.054,大于贪婪搜索得到的输出序列的条件概率。因此,贪婪搜索得到的输出序列“A”“B”“C”“<eos>”并非最优输出序列。
2. 束搜索
束搜索是对贪婪搜索的一个改进算法。它有一个束宽(beam size)超参数。我们将它设为𝑘。在时间步1时,选取当前时间步条件概率最大的𝑘个词,分别组成𝑘个候选输出序列的首词。在之后的每个时间步,基于上个时间步的𝑘个候选输出序列,从𝑘|Y|个可能的输出序列中选取条件概率最大的𝑘个,作为该时间步的候选输出序列。最终,我们从各个时间步的候选输出序列中筛选出包含特殊符号“<eos>”的序列,并将它们中所有特殊符号“<eos>”后面的子序列舍弃,得到最终候选输出序列的集合。
束宽为2,输出序列最大长度为3。候选输出序列有A、C、AB、CE、ABD和CED。
假设输出序列的词典中只包含5个元素,即𝑌={𝐴,𝐵,𝐶,𝐷,𝐸},且其中一个为特殊符号“<eos>”。设束搜索的束宽等于2,输出序列最大长度为3。在输出序列的时间步1时,假设条件概率𝑃(𝑦1∣𝑐)最大的2个词为𝐴和𝐶。在时间步2时将对所有的𝑦2∈𝑌都分别计算𝑃(𝑦2∣𝐴,𝑐)和𝑃(𝑦2∣𝐶,𝑐),并从计算出的10个条件概率中取最大的2个,假设为𝑃(𝐵∣𝐴,𝑐)和𝑃(𝐸∣𝐶,𝑐)。那么,我们在时间步3时将对所有的𝑦3∈𝑌都分别计算𝑃(𝑦3∣𝐴,𝐵,𝑐)和𝑃(𝑦3∣𝐶,𝐸,𝑐),并从计算出的10个条件概率中取最大的2个,假设为𝑃(𝐷∣𝐴,𝐵,𝑐)和𝑃(𝐷∣𝐶,𝐸,𝑐)。如此一来,我们得到6个候选输出序列:(1)𝐴;(2)𝐶;(3)𝐴、𝐵;(4)𝐶、𝐸;(5)𝐴、𝐵、𝐷和(6)𝐶、𝐸、𝐷。接下来,我们将根据这6个序列得出最终候选输出序列的集合。
在最终候选输出序列的集合中,我们取以下分数最高的序列作为输出序列:
(1/𝐿𝛼)*log𝑃(𝑦1,…,𝑦𝐿)=(1/𝐿𝛼)∑(𝑡′=1→𝐿) log𝑃(𝑦𝑡′∣𝑦1,…,𝑦𝑡−1,𝑐)
其中𝐿为最终候选序列长度,𝛼一般可选为0.75。分母上的𝐿𝛼是为了惩罚较长序列在以上分数中较多的对数相加项。分析可知,束搜索的计算开销为𝑂(k|𝑌|𝑇′)。这介于贪婪搜索和穷举搜索的计算开销之间。此外,贪婪搜索可看作是束宽为1的束搜索。束搜索通过灵活的束宽𝑘𝑘来权衡计算开销和搜索质量。
四、注意力机制
1. 计算背景变量
在编码器-解码器架构中,背景变量(context vector)通过注意力机制在每个解码器时间步计算得到。下图描绘了在解码器时间步 t′ 计算背景变量的过程。注意力机制通过函数 a 根据解码器在时间步 t′−1 的隐藏状态和编码器在各个时间步的隐藏状态计算softmax运算的输入。softmax运算输出一个概率分布,并对编码器各个时间步的隐藏状态进行加权平均,从而得到背景变量。
具体计算步骤:
1. 定义隐藏状态:设编码器在时间步 t 的隐藏状态为 ht,总时间步数为 T。
2. 计算背景变量:解码器在时间步 t′ 的背景变量 ct′ 为所有编码器隐藏状态的加权平均:
其中,权重 αt′t 是一个概率分布,表示在时间步 t′ 时编码器隐藏状态 ht 的重要性。
3. 计算权重 αt′t:为得权重 αt′t,使用softmax运算:
4. 计算得分 et′t:得分 et′t 取决于解码器的时间步 t′ 和编码器的时间步 t。它由解码器在时间步 t′−1 的隐藏状态 st′−1 和编码器在时间步 t 的隐藏状态 ht 通过函数 a 计算得到:
2. 矢量化计算
注意力机制的输入包括查询项以及一一对应的键项和值项,其中值项是需要加权平均的一组项。在加权平均中,值项的权重来自查询项以及与该值项对应的键项的计算。
考虑一简单情形,即编码器和解码器的隐藏单元个数均为ℎ,且函数𝑎(𝑠,ℎ)=𝑠^⊤*ℎ。假设我们希望根据解码器单个隐藏状态𝑠𝑡′−1∈𝑅^ℎ和编码器所有隐藏状态ℎ𝑡∈𝑅^ℎ,𝑡=1,…,𝑇来计算背景向量𝑐𝑡′∈𝑅^ℎ。 我们可以将查询项矩阵𝑄∈𝑅^1×ℎ设为𝑠𝑡′−1^⊤,并令键项矩阵𝐾∈𝑅^𝑇×ℎ和值项矩阵𝑉∈𝑅^𝑇×ℎ相同且第𝑡𝑡行均为ℎ𝑡^⊤。此时,只需通过矢量化计算ssoftmax(𝑄𝐾^⊤)*𝑉
即可算出转置后的背景向量𝑐𝑡′^⊤。当查询项矩阵𝑄的行数为𝑛时,上式将得到𝑛行的输出矩阵。输出矩阵与查询项矩阵在相同行上一一对应。
3. 更新隐藏状态
在解码器中,我们可以基于门控循环单元(GRU)的设计稍作修改,以整合注意力机制的背景变量,来更新隐藏状态。具体地,在时间步 t′,解码器的隐藏状态 st′ 可以通过以下公式计算:
其中,⊙ 表示逐元素乘法,zt′ 是更新门,s^~t′ 是候选隐藏状态。
其中,σ 表示 sigmoid 函数,tanh 表示双曲正切函数,W 和 b 是模型的可学习参数。
五、机器翻译
1. 数据预处理
定义一些特殊符号。其中“<pad>”(padding)符号用来添加在较短序列后,直到每个序列等长,而“<bos>”和“<eos>”符号分别表示序列的开始和结束。接着定义两个辅助函数对后面读取的数据进行预处理。
!tar -xf d2lzh_pytorch.tar
import collections
import os
import io
import math
import torch
from torch import nn
import torch.nn.functional as F
import torchtext.vocab as Vocab
import torch.utils.data as Data
import sys
# sys.path.append("..")
import d2lzh_pytorch as d2l
PAD, BOS, EOS = '<pad>', '<bos>', '<eos>'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(torch.__version__, device)
# 将一个序列中所有的词记录在all_tokens中以便之后构造词典,然后在该序列后面添加PAD直到序列
# 长度变为max_seq_len,然后将序列保存在all_seqs中
def process_one_seq(seq_tokens, all_tokens, all_seqs, max_seq_len):
all_tokens.extend(seq_tokens)
seq_tokens += [EOS] + [PAD] * (max_seq_len - len(seq_tokens) - 1)
all_seqs.append(seq_tokens)
# 使用所有的词来构造词典。并将所有序列中的词变换为词索引后构造Tensor
def build_data(all_tokens, all_seqs):
vocab = Vocab.Vocab(collections.Counter(all_tokens),
specials=[PAD, BOS, EOS])
indices = [[vocab.stoi[w] for w in seq] for seq in all_seqs]
return vocab, torch.tensor(indices)
这里使用一个很小的法语—英语数据集。在这个数据集里,每一行是一对法语句子和它对应的英语句子,中间使用'\t'
隔开。在读取数据时,我们在句末附上“<eos>”符号,并可能通过添加“<pad>”符号使每个序列的长度均为max_seq_len
。为法语词和英语词分别创建词典,法语词的索引和英语词的索引相互独立。将序列的最大长度设成7,然后查看读取到的第一个样本。该样本分别包含法语词索引序列和英语词索引序列。
def read_data(max_seq_len):
# in和out分别是input和output的缩写
in_tokens, out_tokens, in_seqs, out_seqs = [], [], [], []
with io.open('fr-en-small.txt') as f:
lines = f.readlines()
for line in lines:
in_seq, out_seq = line.rstrip().split('\t')
in_seq_tokens, out_seq_tokens = in_seq.split(' '), out_seq.split(' ')
if max(len(in_seq_tokens), len(out_seq_tokens)) > max_seq_len - 1:
continue # 如果加上EOS后长于max_seq_len,则忽略掉此样本
process_one_seq(in_seq_tokens, in_tokens, in_seqs, max_seq_len)
process_one_seq(out_seq_tokens, out_tokens, out_seqs, max_seq_len)
in_vocab, in_data = build_data(in_tokens, in_seqs)
out_vocab, out_data = build_data(out_tokens, out_seqs)
return in_vocab, out_vocab, Data.TensorDataset(in_data, out_data)
max_seq_len = 7
in_vocab, out_vocab, dataset = read_data(max_seq_len)
dataset[0]
2. 注意力机制编码器
编码器中,将输入语言的词索引通过词嵌入层得到词的表征,然后输入到一个多层门控循环单元中。PyTorch的nn.GRU
实例在前向计算后也会分别返回输出和最终时间步的多层隐藏状态。其中的输出指的是最后一层的隐藏层在各个时间步的隐藏状态,并不涉及输出层计算,注意力机制将这些输出作为键项和值项。
创建一个批量大小为4、时间步数为7的小批量序列输入。设门控循环单元的隐藏层个数为2,隐藏单元个数为16。编码器对该输入执行前向计算后返回的输出形状为(时间步数, 批量大小, 隐藏单元个数)。门控循环单元在最终时间步的多层隐藏状态的形状为(隐藏层个数, 批量大小, 隐藏单元个数)。对于门控循环单元来说,`state`就是一个元素,即隐藏状态;如果使用长短期记忆,`state`是一个元组,包含两个元素即隐藏状态和记忆细胞。
class Encoder(nn.Module):
def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
drop_prob=0, **kwargs):
super(Encoder, self).__init__(**kwargs)
self.embedding = nn.Embedding(vocab_size, embed_size)
self.rnn = nn.GRU(embed_size, num_hiddens, num_layers, dropout=drop_prob)
def forward(self, inputs, state):
# 输入形状是(批量大小, 时间步数)。将输出互换样本维和时间步维
embedding = self.embedding(inputs.long()).permute(1, 0, 2) # (seq_len, batch, input_size)
return self.rnn(embedding, state)
def begin_state(self):
return None
encoder = Encoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2)
output, state = encoder(torch.zeros((4, 7)), encoder.begin_state())
output.shape, state.shape # GRU的state是h, 而LSTM的是一个元组(h, c)
注意力机制的输入包括查询项、键项和值项。设编码器和解码器的隐藏单元个数相同。这里的查询项为解码器在上一时间步的隐藏状态,形状为(批量大小, 隐藏单元个数);键项和值项均为编码器在所有时间步的隐藏状态,形状为(时间步数, 批量大小, 隐藏单元个数)。注意力机制返回当前时间步的背景变量,形状为(批量大小, 隐藏单元个数)。
编码器的时间步数为10,批量大小为4,编码器和解码器的隐藏单元个数均为8。注意力机制返回一个小批量的背景向量,每个背景向量的长度等于编码器的隐藏单元个数。输出的形状为(4, 8)。
def attention_model(input_size, attention_size):
model = nn.Sequential(nn.Linear(input_size, attention_size, bias=False),
nn.Tanh(),
nn.Linear(attention_size, 1, bias=False))
return model
def attention_forward(model, enc_states, dec_state):
"""
enc_states: (时间步数, 批量大小, 隐藏单元个数)
dec_state: (批量大小, 隐藏单元个数)
"""
# 将解码器隐藏状态广播到和编码器隐藏状态形状相同后进行连结
dec_states = dec_state.unsqueeze(dim=0).expand_as(enc_states)
enc_and_dec_states = torch.cat((enc_states, dec_states), dim=2)
e = model(enc_and_dec_states) # 形状为(时间步数, 批量大小, 1)
alpha = F.softmax(e, dim=0) # 在时间步维度做softmax运算
return (alpha * enc_states).sum(dim=0) # 返回背景变量
seq_len, batch_size, num_hiddens = 10, 4, 8
model = attention_model(2*num_hiddens, 10)
enc_states = torch.zeros((seq_len, batch_size, num_hiddens))
dec_state = torch.zeros((batch_size, num_hiddens))
attention_forward(model, enc_states, dec_state).shape
直接将编码器在最终时间步的隐藏状态作为解码器的初始隐藏状态。这要求编码器和解码器的循环神经网络使用相同的隐藏层个数和隐藏单元个数。在解码器的前向计算中,注意力机制计算得到当前时间步的背景向量。由于解码器的输入来自输出语言的词索引,将输入通过词嵌入层得到表征,然后和背景向量在特征维连结。将连结后的结果与上一时间步的隐藏状态通过门控循环单元计算出当前时间步的输出与隐藏状态。最后,将输出通过全连接层变换为有关各个输出词的预测,形状为(批量大小, 输出词典大小)。
class Decoder(nn.Module):
def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
attention_size, drop_prob=0):
super(Decoder, self).__init__()
self.embedding = nn.Embedding(vocab_size, embed_size)
self.attention = attention_model(2*num_hiddens, attention_size)
# GRU的输入包含attention输出的c和实际输入, 所以尺寸是 num_hiddens+embed_size
self.rnn = nn.GRU(num_hiddens + embed_size, num_hiddens,
num_layers, dropout=drop_prob)
self.out = nn.Linear(num_hiddens, vocab_size)
def forward(self, cur_input, state, enc_states):
"""
cur_input shape: (batch, )
state shape: (num_layers, batch, num_hiddens)
"""
# 使用注意力机制计算背景向量
c = attention_forward(self.attention, enc_states, state[-1])
# 将嵌入后的输入和背景向量在特征维连结, (批量大小, num_hiddens+embed_size)
input_and_c = torch.cat((self.embedding(cur_input), c), dim=1)
# 为输入和背景向量的连结增加时间步维,时间步个数为1
output, state = self.rnn(input_and_c.unsqueeze(0), state)
# 移除时间步维,输出形状为(批量大小, 输出词典大小)
output = self.out(output).squeeze(dim=0)
return output, state
def begin_state(self, enc_state):
# 直接将编码器最终时间步的隐藏状态作为解码器的初始隐藏状态
return enc_state
3. 模型训练
先实现batch_loss
函数计算一个小批量的损失。解码器在最初时间步的输入是特殊字符BOS
。之后,解码器在某时间步的输入为样本输出序列在上一时间步的词,即强制教学。在训练函数中,我们需要同时迭代编码器和解码器的模型参数。接下来,创建模型实例并设置超参数。然后,我们就可以训练模型了。
def batch_loss(encoder, decoder, X, Y, loss):
# 获取批量的大小
batch_size = X.shape[0]
# 初始化编码器的隐藏状态
enc_state = encoder.begin_state()
# 通过编码器获得编码器的输出和最终的隐藏状态
enc_outputs, enc_state = encoder(X, enc_state)
# 使用编码器的最终隐藏状态来初始化解码器的隐藏状态
dec_state = decoder.begin_state(enc_state)
# 初始化解码器在最初时间步的输入为特殊字符BOS的索引
dec_input = torch.tensor([out_vocab.stoi[BOS]] * batch_size)
# 初始化掩码变量mask,用来忽略填充项PAD的损失,初始全为1
mask, num_not_pad_tokens = torch.ones(batch_size,), 0
# 初始化损失值为0
l = torch.tensor([0.0])
# 遍历输出序列的每个时间步,Y.permute(1,0)使得Y的形状为(seq_len, batch)
for y in Y.permute(1, 0): # Y shape: (batch, seq_len)
# 通过解码器得到当前时间步的输出和新的隐藏状态
dec_output, dec_state = decoder(dec_input, dec_state, enc_outputs)
# 计算当前时间步的损失,并按掩码变量加权求和
l = l + (mask * loss(dec_output, y)).sum()
# 使用当前时间步的实际输出作为下一时间步的输入(强制教学)
dec_input = y
# 更新非填充项的计数
num_not_pad_tokens += mask.sum().item()
# 更新掩码变量,保证一旦遇到EOS,接下来的循环中mask就一直是0
mask = mask * (y != out_vocab.stoi[EOS]).float()
# 返回平均损失(注意这里除以的是非填充项的数量)
return l / num_not_pad_tokens
def train(encoder, decoder, dataset, lr, batch_size, num_epochs):
# 定义编码器和解码器的优化器,使用Adam优化算法
enc_optimizer = torch.optim.Adam(encoder.parameters(), lr=lr)
dec_optimizer = torch.optim.Adam(decoder.parameters(), lr=lr)
# 定义损失函数,这里使用CrossEntropyLoss并且不进行损失值的缩减
loss = nn.CrossEntropyLoss(reduction='none')
# 创建数据加载器,打乱数据顺序并分批次处理
data_iter = Data.DataLoader(dataset, batch_size, shuffle=True)
# 开始训练循环
for epoch in range(num_epochs):
# 初始化本epoch的损失总和
l_sum = 0.0
# 遍历数据加载器中的每一个批次数据
for X, Y in data_iter:
# 梯度清零,防止梯度累积
enc_optimizer.zero_grad()
dec_optimizer.zero_grad()
# 计算当前批次的损失
l = batch_loss(encoder, decoder, X, Y, loss)
# 反向传播计算梯度
l.backward()
# 更新编码器和解码器的参数
enc_optimizer.step()
dec_optimizer.step()
# 累加损失值
l_sum += l.item()
# 每10个epoch输出一次当前的平均损失
if (epoch + 1) % 10 == 0:
print("epoch %d, loss %.3f" % (epoch + 1, l_sum / len(data_iter)))
# 设置模型和训练的超参数
embed_size, num_hiddens, num_layers = 64, 64, 2 # 嵌入维度、隐藏单元数、层数
attention_size, drop_prob, lr, batch_size, num_epochs = 10, 0.5, 0.01, 2, 50
# attention_size: 注意力机制的维度
# drop_prob: dropout的概率
# lr: 学习率
# batch_size: 批次大小
# num_epochs: 训练轮数
# 初始化编码器
encoder = Encoder(len(in_vocab), embed_size, num_hiddens, num_layers, drop_prob)
# Encoder类需要提供输入词汇表的大小、嵌入维度、隐藏单元数、层数以及dropout概率
# 初始化解码器
decoder = Decoder(len(out_vocab), embed_size, num_hiddens, num_layers, attention_size, drop_prob)
# Decoder类需要提供输出词汇表的大小、嵌入维度、隐藏单元数、层数、注意力机制的维度以及dropout概率
# 训练模型
train(encoder, decoder, dataset, lr, batch_size, num_epochs)
# 调用train函数进行训练,传入编码器、解码器、数据集、学习率、批次大小和训练轮数
4. 预测与评价
此实现最简单的贪婪搜索,简单测试模型。输入法语句子“ils regardent.”,翻译后的英语句子应该是“they are watching.”。
BLEU评价机器翻译结果时考虑预测序列中与标签序列匹配的子序列的精度。具体来说,设预测序列长度为 lenpred,标签序列长度为 lenlabel,则精度 pn 表示预测序列中匹配长度为 n 的子序列的数量与预测序列中长度为 n 的子序列总数的比例。例如,如果标签序列为 A,B,C,D,E,F 预测序列为 A,B,B,C,D,则 p1=4/5,p2=3/4,p3=1/3,p4=0。
BLEU的定义为:
其中:k 是考虑的最大子序列长度。wn 是权重,通常为 1/k,用来减少较短输出的影响。
import math
import collections
def bleu(pred_tokens, label_tokens, k):
"""
计算BLEU分数
Args:
- pred_tokens (list): 预测序列的单词列表
- label_tokens (list): 标签序列的单词列表
- k (int): 最大匹配的子序列长度
Returns:
- float: BLEU分数
"""
len_pred, len_label = len(pred_tokens), len(label_tokens)
# 计算惩罚系数,惩罚较短预测序列的长度
penalty = math.exp(min(0, 1 - len_label / len_pred))
# 初始化BLEU分数
score = penalty
# 计算匹配子序列的精度
for n in range(1, k + 1):
num_matches = 0
label_subs = collections.defaultdict(int)
# 构建标签序列中长度为n的所有子序列的字典
for i in range(len_label - n + 1):
label_subs[''.join(label_tokens[i: i + n])] += 1
# 计算预测序列中长度为n的匹配子序列数
for i in range(len_pred - n + 1):
if label_subs[''.join(pred_tokens[i: i + n])] > 0:
num_matches += 1
label_subs[''.join(pred_tokens[i: i + n])] -= 1
# 计算长度为n的匹配子序列的精度
precision_n = num_matches / (len_pred - n + 1)
# 加权乘以惩罚系数,根据子序列长度调整权重
score *= math.pow(precision_n, math.pow(0.5, n))
return score
def score(input_seq, label_seq, k):
# 使用机器翻译模型预测输入序列的翻译结果
pred_tokens = translate(encoder, decoder, input_seq, max_seq_len)
# 将标签序列拆分为单词列表
label_tokens = label_seq.split(' ')
# 计算预测结果与标签序列的BLEU分数
bleu_score = bleu(pred_tokens, label_tokens, k)
# 输出BLEU分数和预测的翻译结果
print('BLEU score: %.3f, Predicted translation: %s' % (bleu_score, ' '.join(pred_tokens)))
六、Transformer
1. 数据预处理
导入所需的包
import math
import torchtext # 文本数据处理模块
import torch
import torch.nn as nn # 神经网络模块
from torch import Tensor # 从PyTorch中导入Tensor类
from torch.nn.utils.rnn import pad_sequence # pad_sequence函数用于填充序列使其具有相同长度
from torch.utils.data import DataLoader # DataLoader类用于批量加载数据
from collections import Counter # Counter类用于统计词频
from torchtext.vocab import Vocab # Vocab类用于创建词汇表
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
import io # 处理输入输出操作模块
import time # 时间操作模块
import pandas as pd # 数据处理模块
import numpy as np
import pickle # pickle模块用于对象序列化
import tqdm # 进度条
import sentencepiece as spm # 文本分词模块
torch.manual_seed(0) # 设置随机种子为0,以确保结果的可重复性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(torch.cuda.get_device_name(0))
获取数据集,在导入所有日语及其英语对应文本后,删除数据集中的最后一条数据,因为它有缺失值。总共,trainen和trainja中的句子数量为5,973,071条。通常建议对数据进行抽样,并确保一切按预期运行,然后再一次性使用所有数据,以节省时间。
import pandas as pd
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)
# trainja.pop(5972)
2. 分词与对象构建
与英语或其他字母语言不同,日语句子中没有空格来分隔单词。使用JParaCrawl提供的分词器,它们使用了SentencePiece来处理日语和英语。
import sentencepiece as spm
# 加载英文的SentencePiece模型,并创建Tokenizer对象
en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')
# 加载日文的SentencePiece模型,并创建Tokenizer对象
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')
使用分词器和原始句子,接着从TorchText中导入Vocab对象。此过程的时间长度取决于数据集的大小和计算能力,可能需要几秒钟到几分钟不等。不同的分词器也会影响构建词汇表所需的时间。
from collections import Counter
from torchtext.vocab import Vocab
def build_vocab(sentences, tokenizer):
"""
根据输入的句子列表构建词汇表。
Args:
- sentences (list): 输入句子的列表。
- tokenizer: 分词器对象,例如使用 SentencePieceProcessor。
Returns:
- Vocab: 构建好的词汇表对象。
"""
counter = Counter()
# 遍历每个句子,并更新计数器
for sentence in sentences:
# 使用分词器对句子进行分词并编码(转换为字符串类型的编码)
counter.update(tokenizer.encode(sentence, out_type=str))
# 使用计数器构建词汇表,同时指定特殊词汇(unknown, padding, beginning-of-sequence, end-of-sequence)
return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
# 使用 build_vocab 函数分别为日语和英语构建词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer) # 构建日语词汇表
en_vocab = build_vocab(trainen, en_tokenizer) # 构建英语词汇表
import torch
def data_process(ja, en):
"""
将原始的日语和英语文本数据处理为模型可以处理的张量格式数据集。
Args:
- ja (list): 原始日语文本数据列表。
- en (list): 原始英语文本数据列表。
Returns:
- list: 包含处理后的数据对 (日语张量, 英语张量) 的列表。
"""
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
# 使用 data_process 函数处理训练数据集 trainja 和 trainen
train_data = data_process(trainja, trainen)
创建用于训练迭代的DataLoader对象,将BATCH_SIZE设置为16,以防止“cuda内存不足”的问题,但这取决于各种因素,如计算机内存容量、数据大小等等。
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
BATCH_SIZE = 8 # 批处理大小
PAD_IDX = ja_vocab['<pad>'] # 填充索引
BOS_IDX = ja_vocab['<bos>'] # 序列开始索引
EOS_IDX = ja_vocab['<eos>'] # 序列结束索引
def generate_batch(data_batch):
"""
生成一个批次的数据,并进行填充以保证批次中的张量长度一致。
Args:
- data_batch (list): 包含批次数据对的列表,每个数据对为 (日语张量, 英语张量)。
Returns:
- tuple: 包含填充后的日语批次张量和英语批次张量的元组。
"""
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))
# 在英语张量的开头添加 <bos> 标记,在末尾添加 <eos> 标记
en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
# 对日语和英语批次张量进行填充,使用 <pad> 标记进行填充
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_data 进行批处理
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
shuffle=True, collate_fn=generate_batch)
3. Sequence-to-sequence
Transformer是一种Seq2Seq模型,用于解决机器翻译任务。Transformer模型包括一个编码器和一个解码器模块,每个模块包含固定数量的层。编码器通过一系列的多头注意力和前馈网络层处理输入序列。编码器的输出称为记忆(memory),将与目标张量一起馈送给解码器。编码器和解码器使用教师强制(teacher forcing)技术进行端到端的训练。
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
NHEAD = 8
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, tgt: 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(tgt))
# 编码器阶段,生成编码器的记忆
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)
文本标记通过使用标记嵌入来表示。为了引入单词顺序的概念,还会添加位置编码到标记嵌入中。
import torch
import torch.nn as nn
import math
from torch import Tensor
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)
# 使用 Dropout 对位置编码进行随机失活
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__()
# 使用 nn.Embedding 创建词嵌入层
self.embedding = nn.Embedding(vocab_size, emb_size)
self.emb_size = emb_size
def forward(self, tokens: Tensor):
# 对输入的 tokens 进行词嵌入操作,并乘以 sqrt(emb_size) 进行缩放
return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
创建一个后续词掩码,用于阻止目标词关注后续词,还创建了用于遮蔽源和目标填充标记的掩码。
def generate_square_subsequent_mask(sz):
"""
生成一个正方形的下三角掩码矩阵,用于解码器中的自注意力机制。
Args:
- sz (int): 矩阵的大小,即序列长度。
Returns:
- torch.Tensor: 生成的下三角掩码矩阵,形状为 (sz, sz)。
"""
# 创建一个全为 1 的矩阵,对角线及其以上为 1,其余为 0
mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
# 将掩码转换为浮点型,并将值为 0 的位置用负无穷大来替换
mask = mask.float().masked_fill(mask == 0, float('-inf'))
# 将值为 1 的位置用 0.0 来替换
mask = mask.masked_fill(mask == 1, float(0.0))
return mask
def create_mask(src, tgt):
"""
生成用于 Transformer 模型的掩码矩阵。
Args:
- src (torch.Tensor): 源语言张量,形状为 (src_seq_len, batch_size).
- tgt (torch.Tensor): 目标语言张量,形状为 (tgt_seq_len, batch_size).
Returns:
- torch.Tensor: 源语言掩码矩阵,形状为 (src_seq_len, src_seq_len).
- torch.Tensor: 目标语言掩码矩阵,形状为 (tgt_seq_len, tgt_seq_len).
- torch.Tensor: 源语言填充掩码,形状为 (batch_size, src_seq_len).
- torch.Tensor: 目标语言填充掩码,形状为 (batch_size, tgt_seq_len).
"""
src_seq_len = src.shape[0]
tgt_seq_len = tgt.shape[0]
# 生成目标语言的下三角掩码矩阵
tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
# 源语言掩码矩阵,全为 False(即全为 0)
src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
# 源语言填充掩码,标记为 True(即为 1)的位置是 PAD_IDX
src_padding_mask = (src == PAD_IDX).transpose(0, 1)
# 目标语言填充掩码,标记为 True(即为 1)的位置是 PAD_IDX
tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask
# 导入所需的库和模块
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence
from torch import Tensor
import math
# 设置一些超参数和常量
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)
# 初始化模型参数(使用 Xavier 初始化)
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)
# 定义优化器(Adam 优化器)
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 # 初始化损失值为 0
for idx, (src, tgt) in enumerate(train_iter):
src = src.to(device) # 将源语言数据移动到指定的设备(GPU 或 CPU)
tgt = tgt.to(device) # 将目标语言数据移动到指定的设备(GPU 或 CPU)
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 # 初始化损失值为 0
with torch.no_grad(): # 在评估阶段不计算梯度
for idx, (src, tgt) in enumerate(val_iter):
src = src.to(device) # 将源语言数据移动到指定的设备(GPU 或 CPU)
tgt = tgt.to(device) # 将目标语言数据移动到指定的设备(GPU 或 CPU)
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) # 返回平均损失
4. 模型训练
import tqdm
import time
# 使用 tqdm 显示训练进度
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 的结束时间
# 打印当前 epoch 的训练损失和训练时间
print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Epoch time = {(end_time - start_time):.3f}s")
def greedy_decode(model, src, src_mask, max_len, start_symbol):
"""
使用贪婪解码方法生成目标语言序列。
Args:
- model: Transformer 模型
- src: 源语言序列张量
- src_mask: 源语言序列的掩码张量
- max_len: 生成序列的最大长度
- start_symbol: 目标语言序列的起始符号
Returns:
- ys: 生成的目标语言序列张量
"""
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):
"""
使用训练好的 Transformer 模型进行翻译。
Args:
- model: 训练好的 Seq2SeqTransformer 模型
- src: 源语言文本字符串
- src_vocab: 源语言词汇表
- tgt_vocab: 目标语言词汇表
- src_tokenizer: 源语言文本的分词器
Returns:
- translation: 翻译后的目标语言文本字符串
"""
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() # 使用贪婪解码生成目标语言序列
translation = " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "") # 将目标语言序列转换为文本字符串
return translation
translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)",
ja_vocab, en_vocab, ja_tokenizer)
translate(transformer, "(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)
trainen.pop(5)
trainja.pop(5)
5. 保存模型对象
首先使用Pickle保存Vocab对象(en_vocab 和 ja_vocab)。使用PyTorch的保存和加载函数来保存模型以供以后使用。
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')