背景
最近对序列模型比较感兴趣,正好使用pytorch实现一个简单的机器翻译模型,在gpu运行环境下,在epoch=50轮左右的时候,训练效果,看似还不错。这里记录下。
效果
- 训练后的效果如下
- 正常情况下,翻译的结果和正确答案差异不大。
- 但是在某些特殊情况下,比如字符在测试集合中有,但是训练集合中没有的情况下,效果会比较差。
代码实现
数据预处理
import numpy as np
import unicodedata
from torch.autograd import Variable
import torch
import re
from collections import Counter
UNK = 0 # 未登录词的标识符对应的词典id
PAD = 1 # padding占位符对应的词典id
import os
# os.environ["CUDA_VISIBLE_DEVICES"] = '2'
# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def seq_padding(X, padding=0):
"""
对一个batch批次(以单词id表示)的数据进行padding填充对齐长度
"""
# 计算该批次数据各条数据句子长度
L = [len(x) for x in X]
# 获取该批次数据最大句子长度
ML = max(L)
# 对X中各条数据x进行遍历,如果长度短于该批次数据最大长度ML,则以padding id填充缺失长度ML-len(x)
return np.array([
np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X
])
def subsequent_mask(size):
"Mask out subsequent positions."
# 设定subsequent_mask矩阵的shape
attn_shape = (1, size, size)
# TODO: 生成一个右上角(不含主对角线)为全1,左下角(含主对角线)为全0的subsequent_mask矩阵
subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
# TODO: 返回一个右上角(不含主对角线)为全False,左下角(含主对角线)为全True的subsequent_mask矩阵
return torch.from_numpy(subsequent_mask) == 0
class Batch:
"Object for holding a batch of data with mask during training."
def __init__(self, src, trg=None, pad=0, DEVICE=None):
# 将输入与输出的单词id表示的数据规范成整数类型
src = torch.from_numpy(src).to(DEVICE).long()
trg = torch.from_numpy(trg).to(DEVICE).long()
self.src = src
# 对于当前输入的句子非空部分进行判断成bool序列
# 并在seq length前面增加一维,形成维度为 1×seq length 的矩阵
self.src_mask = (src != pad).unsqueeze(-2).to(DEVICE)
# 如果输出目标不为空,则需要对decoder要使用到的target句子进行mask
if trg is not None:
# decoder要用到的target输入部分
self.trg = trg[:, :-1]
# decoder训练时应预测输出的target结果
self.trg_y = trg[:, 1:]
# 将target输入部分进行attention mask
self.trg_mask = self.make_std_mask(self.trg, pad).to(DEVICE)
# 将应输出的target结果中实际的词数进行统计
self.ntokens = (self.trg_y != pad).data.sum().to(DEVICE)
# Mask掩码操作
@staticmethod
def make_std_mask(tgt, pad):
"Create a mask to hide padding and future words."
tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & Variable(subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
return tgt_mask
class PrepareData:
def __init__(self, train_file, dev_file, batch_size, DEVICE):
self.batch_size = batch_size
# 读取数据 并分词
self.train_en, self.train_cn = self.load_data(train_file)
self.dev_en, self.dev_cn = self.load_data(dev_file)
# 构建单词表
self.en_word_dict, self.en_total_words, self.en_index_dict = self.build_dict(self.train_en)
self.cn_word_dict, self.cn_total_words, self.cn_index_dict = self.build_dict(self.train_cn)
self.tgt_vocab = len(self.en_word_dict)
self.src_vocab = len(self.cn_word_dict)
# id化
self.train_en, self.train_cn = self.wordToID(self.train_en, self.train_cn, self.en_word_dict, self.cn_word_dict)
self.dev_en, self.dev_cn = self.wordToID(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)
# 划分batch + padding + mask
self.train_data = self.splitBatch(self.train_en, self.train_cn, batch_size, DEVICE)
self.dev_data = self.splitBatch(self.dev_en, self.dev_cn, batch_size, DEVICE)
def unicodeToAscii(self, s):
return ''.join(
c for c in unicodedata.normalize('NFD', s)
if unicodedata.category(c) != 'Mn'
)
def normalizeString(self, s):
s = self.unicodeToAscii(s.lower().strip())
s = re.sub(r"([.!?])", r" \1", s)
# s = re.sub(r"[^a-zA-Z.!?]+", r" ", s)
return s
def load_data(self, path):
"""
读取翻译前(英文)和翻译后(中文)的数据文件
每条数据都进行分词,然后构建成包含起始符(BOS)和终止符(EOS)的单词(中文为字符)列表
形式如:en = [['BOS', 'i', 'love', 'you', 'EOS'], ['BOS', 'me', 'too', 'EOS'], ...]
cn = [['BOS', '我', '爱', '你', 'EOS'], ['BOS', '我', '也', '是', 'EOS'], ...]
"""
en = []
cn = []
# TODO ...
with open(path, 'r', encoding='utf-8') as fin:
for line in fin:
list_content = line.split('\t')
# print(list_content)
# engs = ['BOS'] + word_tokenize(list_content[0]) + ['EOS']
engs = ['BOS'] + self.normalizeString(list_content[0]).split(' ') + ['EOS']
# print('engs=', engs)
en.append(engs)
# print('list_content[1]=', list_content[1])
# lists = word_tokenize(" ".join(list_content[1]))
lists = words = [w for w in self.normalizeString(list_content[1])]
# print('lists=', lists)
chins = ['BOS'] + lists + ['EOS']
# print('chins=', chins)
cn.append(chins)
return en, cn
def build_dict(self, sentences, max_words=50000):
"""
传入load_data构造的分词后的列表数据
构建词典(key为单词,value为id值)
"""
# 对数据中所有单词进行计数
# print('sentences=', sentences)
word_count = Counter()
for sentence in sentences:
for s in sentence:
word_count[s] += 1
# 只保留最高频的前max_words数的单词构建词典
# 并添加上UNK和PAD两个单词,对应id已经初始化设置过
ls = word_count.most_common(max_words)
# 统计词典的总词数
total_words = len(ls) + 2
word_dict = {w[0]: index + 2 for index, w in enumerate(ls)}
word_dict['UNK'] = UNK
word_dict['PAD'] = PAD
# 再构建一个反向的词典,供id转单词使用
index_dict = {v: k for k, v in word_dict.items()}
return word_dict, total_words, index_dict
def wordToID(self, en, cn, en_dict, cn_dict, sort=True):
"""
该方法可以将翻译前(英文)数据和翻译后(中文)数据的单词列表表示的数据
均转为id列表表示的数据
如果sort参数设置为True,则会以翻译前(英文)的句子(单词数)长度排序
以便后续分batch做padding时,同批次各句子需要padding的长度相近减少padding量
"""
# 计算英文数据条数
length = len(en)
# TODO: 将翻译前(英文)数据和翻译后(中文)数据都转换为id表示的形式
out_en_ids = [[en_dict.get(w, 0) for w in sent] for sent in en]
out_cn_ids = [[cn_dict.get(w, 0) for w in sent] for sent in cn]
# 构建一个按照句子长度排序的函数
def len_argsort(seq):
"""
传入一系列句子数据(分好词的列表形式),
按照句子长度排序后,返回排序后原来各句子在数据中的索引下标
"""
return sorted(range(len(seq)), key=lambda x: len(seq[x]))
# 把中文和英文按照同样的顺序排序
if sort:
# 以英文句子长度排序的(句子下标)顺序为基准
sorted_index = len_argsort(out_en_ids)
# TODO: 对翻译前(英文)数据和翻译后(中文)数据都按此基准进行排序
out_en_ids = [out_en_ids[i] for i in sorted_index]
out_cn_ids = [out_cn_ids[i] for i in sorted_index]
return out_en_ids, out_cn_ids
def splitBatch(self, en, cn, batch_size, shuffle=True, DEVICE=None):
"""
将以单词id列表表示的翻译前(英文)数据和翻译后(中文)数据
按照指定的batch_size进行划分
如果shuffle参数为True,则会对这些batch数据顺序进行随机打乱
"""
# 在按数据长度生成的各条数据下标列表[0, 1, ..., len(en)-1]中
# 每隔指定长度(batch_size)取一个下标作为后续生成batch的起始下标
idx_list = np.arange(0, len(en), batch_size)
# 如果shuffle参数为True,则将这些各batch起始下标打乱
if shuffle:
np.random.shuffle(idx_list)
# 存放各个batch批次的句子数据索引下标
batch_indexs = []
for idx in idx_list:
# 注意,起始下标最大的那个batch可能会超出数据大小
# 因此要限定其终止下标不能超过数据大小
batch_indexs.append(np.arange(idx, min(idx + batch_size, len(en))))
# 按各batch批次的句子数据索引下标,构建实际的单词id列表表示的各batch句子数据
batches = []
for batch_index in batch_indexs:
# 按当前batch的各句子下标(数组批量索引)提取对应的单词id列表句子表示数据
batch_en = [en[index] for index in batch_index]
batch_cn = [cn[index] for index in batch_index]
# 对当前batch的各个句子都进行padding对齐长度
# 维度为:batch数量×batch_size×每个batch最大句子长度
batch_cn = seq_padding(batch_cn)
batch_en = seq_padding(batch_en)
# 将当前batch的英文和中文数据添加到存放所有batch数据的列表中
batches.append(Batch(batch_cn, batch_en, DEVICE=DEVICE))
return batches
序列化模型实现
import torch.nn as nn
import torch
import copy
import torch.nn.functional as F
from torch.autograd import Variable
import math
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def clones(module, N):
"""
克隆模型块,克隆的模型块参数不共享
"""
return nn.ModuleList([copy.deepcopy(module).to(DEVICE) for _ in range(N)])
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
super(MultiHeadedAttention, self).__init__()
# 保证可以整除
assert d_model % h == 0
# 得到一个head的attention表示维度
self.d_k = d_model // h
# head数量
self.h = h
# 定义4个全连接函数,供后续作为WQ,WK,WV矩阵和最后h个多头注意力矩阵concat之后进行变换的矩阵
self.linears = clones(nn.Linear(d_model, d_model), 4)
self.attn = None
self.dropout = nn.Dropout(p=dropout)
self.to(DEVICE)
def forward(self, query, key, value, mask=None):
if mask is not None:
mask = mask.unsqueeze(1)
# query的第一个维度值为batch size
nbatches = query.size(0)
# print('nbatches=', nbatches)
# 将embedding层乘以WQ,WK,WV矩阵(均为全连接)
# 并将结果拆成h块,然后将第二个和第三个维度值互换(具体过程见上述解析)
query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linears, (query, key, value))]
# 调用上述定义的attention函数计算得到h个注意力矩阵跟value的乘积,以及注意力矩阵
x, self.attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)
# 将h个多头注意力矩阵concat起来(注意要先把h变回到第三维的位置)
x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
# 使用self.linears中构造的最后一个全连接函数来存放变换后的矩阵进行返回
return self.linears[-1](x)
def attention(self, query, key, value, mask=None, dropout=None):
# 将query矩阵的最后一个维度值作为d_k
d_k = query.size(-1)
# TODO: 将key的最后两个维度互换(转置),才能与query矩阵相乘,乘完了还要除以d_k开根号
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
# 如果存在要进行mask的内容,则将那些为0的部分替换成一个很大的负数
scores = scores.to(DEVICE)
mask = mask.to(DEVICE)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# TODO: 将mask后的attention矩阵按照最后一个维度进行softmax
p_attn = F.softmax(scores, dim=-1)
# 如果dropout参数设置为非空,则进行dropout操作
if dropout is not None:
p_attn = dropout(p_attn)
# 最后返回注意力矩阵跟value的乘积,以及注意力矩阵
return torch.matmul(p_attn, value), p_attn
class PositionwiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super(PositionwiseFeedForward, self).__init__()
self.w_1 = nn.Linear(d_model, d_ff)
self.w_2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
self.to(DEVICE)
def forward(self, x):
# TODO: 请利用init中的成员变量实现Feed Forward层的功能
return self.w_2(self.dropout(F.relu(self.w_1(x))))
class LayerNorm(nn.Module):
def __init__(self, features, eps=1e-6):
super(LayerNorm, self).__init__()
# 初始化α为全1, 而β为全0
self.a_2 = nn.Parameter(torch.ones(features))
self.b_2 = nn.Parameter(torch.zeros(features))
# 平滑项
self.eps = eps
self.to(DEVICE)
def forward(self, x):
# TODO: 请利用init中的成员变量实现LayerNorm层的功能
# 按最后一个维度计算均值和方差
mean = x.mean(-1, keepdim=True)
std = x.std(-1, keepdim=True)
# TODO: 返回Layer Norm的结果
return self.a_2 * (x - mean) / torch.sqrt(std ** 2 + self.eps) + self.b_2
class SublayerConnection(nn.Module):
"""
SublayerConnection的作用就是把Multi-Head Attention和Feed Forward层连在一起
只不过每一层输出之后都要先做Layer Norm再残差连接
sublayer是lamda函数
"""
def __init__(self, size, dropout):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(size)
self.dropout = nn.Dropout(dropout)
self.to(DEVICE)
def forward(self, x, sublayer):
# TODO: 请利用init中的成员变量实现LayerNorm和残差连接的功能
# 返回Layer Norm和残差连接后结果
return x + self.dropout(sublayer(self.norm(x)))
class EncoderLayer(nn.Module):
def __init__(self, h, d_model, d_ff, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = MultiHeadedAttention(h, d_model).to(DEVICE)
self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)
# SublayerConnection的作用就是把multi和ffn连在一起
# 只不过每一层输出之后都要先做Layer Norm再残差连接
self.sublayer = clones(SublayerConnection(d_model, dropout), 2)
# d_model
self.size = d_model
self.to(DEVICE)
def forward(self, x, mask):
# 将embedding层进行Multi head Attention
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
# 注意到attn得到的结果x直接作为了下一层的输入
return self.sublayer[1](x, self.feed_forward)
class Encoder(nn.Module):
# layer = EncoderLayer
# N = 6
def __init__(self, h, d_model, d_ff, dropout, N):
super(Encoder, self).__init__()
# 复制N个encoder layer
layer = EncoderLayer(h, d_model, d_ff, dropout).to(DEVICE)
self.layers = clones(layer, N)
# Layer Norm
self.norm = LayerNorm(d_model).to(DEVICE)
self.to(DEVICE)
def forward(self, x, mask):
"""
使用循环连续eecode N次(这里为6次)
这里的Eecoderlayer会接收一个对于输入的attention mask处理
"""
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
class DecoderLayer(nn.Module):
def __init__(self, h, d_model, d_ff, dropout):
super(DecoderLayer, self).__init__()
self.size = d_model
# Self-Attention
self.self_attn = MultiHeadedAttention(h, d_model).to(DEVICE)
# 与Encoder传入的Context进行Attention
self.src_attn = MultiHeadedAttention(h, d_model).to(DEVICE)
self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
self.sublayer = clones(SublayerConnection(d_model, dropout), 3)
self.to(DEVICE)
def forward(self, x, memory, src_mask, tgt_mask):
# 用m来存放encoder的最终hidden表示结果
m = memory
# TODO: 参照EncoderLayer完成DecoderLayer的forwark函数
# Self-Attention:注意self-attention的q,k和v均为decoder hidden
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
# Context-Attention:注意context-attention的q为decoder hidden,而k和v为encoder hidden
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
return self.sublayer[2](x, self.feed_forward)
class Decoder(nn.Module):
def __init__(self, h, d_model, d_ff, dropout, N, vocab):
super(Decoder, self).__init__()
layer = DecoderLayer(h, d_model, d_ff, dropout)
# 复制N个encoder layer
self.layers = clones(layer, N)
# Layer Norm
self.norm = LayerNorm(layer.size)
self.proj = nn.Linear(d_model, vocab)
self.to(DEVICE)
def forward(self, x, memory, src_mask, tgt_mask, voc_out=False):
"""
使用循环连续decode N次(这里为6次)
这里的Decoderlayer会接收一个对于输入的attention mask处理
和一个对输出的attention mask + subsequent mask处理
"""
for layer in self.layers:
x = layer(x, memory, src_mask, tgt_mask)
x = self.norm(x)
if not voc_out:
return x
else:
return F.log_softmax(self.proj(x), dim=-1)
class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout, DEVICE, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
# 初始化一个size为 max_len(设定的最大长度)×embedding维度 的全零矩阵
# 来存放所有小于这个长度位置对应的porisional embedding
pe = torch.zeros(max_len, d_model, device=DEVICE)
# 生成一个位置下标的tensor矩阵(每一行都是一个位置下标)
"""
形式如:
tensor([[0.],
[1.],
[2.],
[3.],
[4.],
...])
"""
position = torch.arange(0., max_len, device=DEVICE).unsqueeze(1)
# 这里幂运算太多,我们使用exp和log来转换实现公式中pos下面要除以的分母(由于是分母,要注意带负号)
div_term = torch.exp(torch.arange(0., d_model, 2, device=DEVICE) * -(math.log(10000.0) / d_model))
# TODO: 根据公式,计算各个位置在各embedding维度上的位置纹理值,存放到pe矩阵中
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
# 加1个维度,使得pe维度变为:1×max_len×embedding维度
# (方便后续与一个batch的句子所有词的embedding批量相加)
pe = pe.unsqueeze(0)
# 将pe矩阵以持久的buffer状态存下(不会作为要训练的参数)
self.register_buffer('pe', pe)
self.to(DEVICE)
def forward(self, x):
# 将一个batch的句子所有词的embedding与已构建好的positional embeding相加
# (这里按照该批次数据的最大句子长度来取对应需要的那些positional embedding值)
x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
return self.dropout(x)
class Embeddings(nn.Module):
def __init__(self, d_model, vocab, DEVICE):
super(Embeddings, self).__init__()
# Embedding层
self.lut = nn.Embedding(vocab, d_model).to(DEVICE)
# Embedding维数
self.d_model = d_model
self.DEVICE = DEVICE
self.to(DEVICE)
def forward(self, x):
# 返回x对应的embedding矩阵(需要乘以math.sqrt(d_model))
x = x.to(self.DEVICE)
return self.lut(x) * math.sqrt(self.d_model)
# class Generator(nn.Module):
# # vocab: tgt_vocab
# def __init__(self, d_model, vocab):
# super(Generator, self).__init__()
# # decode后的结果,先进入一个全连接层变为词典大小的向量
# self.proj = nn.Linear(d_model, vocab)
#
# def forward(self, x):
# # 然后再进行log_softmax操作(在softmax结果上再做多一次log运算)
# return F.log_softmax(self.proj(x), dim=-1)
class Transformer(nn.Module):
def __init__(self, h, d_model, d_ff, dropout, N, src_vocab, tgt_vocab, DEVICE):
super(Transformer, self).__init__()
self.encoder = Encoder(h, d_model, d_ff, dropout, N).to(DEVICE)
self.decoder = Decoder(h, d_model, d_ff, dropout, N, tgt_vocab).to(DEVICE)
self.src_embed = nn.Sequential(Embeddings(d_model, src_vocab, DEVICE).to(DEVICE), PositionalEncoding(d_model, dropout, DEVICE).to(DEVICE))
self.tgt_embed = nn.Sequential(Embeddings(d_model, tgt_vocab, DEVICE).to(DEVICE), PositionalEncoding(d_model, dropout, DEVICE).to(DEVICE))
# self.generator = Generator(d_model, tgt_vocab)
for p in self.parameters():
if p.dim() > 1:
# 这里初始化采用的是nn.init.xavier_uniform
nn.init.xavier_uniform_(p)
self.to(DEVICE)
def encode(self, src, src_mask):
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask, voc_out=True)
def forward(self, src, tgt, src_mask, tgt_mask):
# encoder的结果作为decoder的memory参数传入,进行decode
return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
模型训练
import time
import torch
import torch.nn as nn
from torch.autograd import Variable
import os
import random
from data_generator import *
import numpy as np
from transformer_models import *
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class CustomLearnRateOptimizer:
"Optim wrapper that implements rate."
def __init__(self, model_size, factor, warmup, model):
self.optimizer = torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)
self._step = 0
self.warmup = warmup
self.factor = factor
self.model_size = model_size
self._rate = 0
def step(self):
"Update parameters and rate"
self._step += 1
rate = self.rate()
for p in self.optimizer.param_groups:
p['lr'] = rate
self._rate = rate
self.optimizer.step()
def rate(self, step=None):
"Implement `lrate` above"
if step is None:
step = self._step
return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))
class LabelSmoothingLoss(nn.Module):
"""标签平滑处理"""
def __init__(self, size, padding_idx=0, smoothing=0.0):
super(LabelSmoothingLoss, self).__init__()
self.criterion = nn.KLDivLoss(reduction='sum')
self.padding_idx = padding_idx
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
self.size = size
self.true_dist = None
def forward(self, x, target):
assert x.size(1) == self.size
x = x.to(DEVICE)
target = target.to(DEVICE)
true_dist = x.data.clone()
true_dist.fill_(self.smoothing / (self.size - 2))
true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
true_dist[:, self.padding_idx] = 0
mask = torch.nonzero(target.data == self.padding_idx)
if mask.dim() > 0:
true_dist.index_fill_(0, mask.squeeze(), 0.0)
self.true_dist = true_dist
return self.criterion(x, Variable(true_dist, requires_grad=False))
class Trainer:
def __init__(self, epochs, tgt_vocab, D_MODEL, model, SAVE_FILE='model1.pt', MAX_LENGTH=100):
self.epochs = epochs
self.loss = LabelSmoothingLoss(tgt_vocab)
self.optimizer = CustomLearnRateOptimizer(D_MODEL, 1, 2000, model)
self.SAVE_FILE = SAVE_FILE
self.MAX_LENGTH = MAX_LENGTH
self.model = model
def compute_loss(self, x, y, norm):
# 计算loss,并且反向传播
loss = self.loss(x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)) / norm
# print('loss=', loss.shape, loss)
loss.backward()
if self.optimizer is not None:
self.optimizer.step()
self.optimizer.optimizer.zero_grad()
return loss.data.item() * norm.float()
def train(self, data):
"""
训练并保存模型
"""
for p in self.model.parameters():
if p.dim() > 1:
# 这里初始化采用的是nn.init.xavier_uniform
nn.init.xavier_uniform_(p)
# 初始化模型在dev集上的最优Loss为一个较大值
best_dev_loss = 1e5
for epoch in range(self.epochs):
# 模型训练
self.model.train()
self.run_epoch(data.train_data, self.model, epoch)
self.model.eval()
# 在dev集上进行loss评估
print('>>>>> Evaluate')
dev_loss = self.run_epoch(data.dev_data, self.model, epoch, max_step=200)
print('<<<<< Evaluate loss: %f' % dev_loss)
# # TODO: 如果当前epoch的模型在dev集上的loss优于之前记录的最优loss则保存当前模型,并更新最优loss值
if dev_loss < best_dev_loss:
print('saving model...')
torch.save(self.model.state_dict(), self.SAVE_FILE)
best_dev_loss = dev_loss
print('>>>>>>>>>>>>>>> Evaluate case, epach=%d' % epoch)
self.evaluate(data)
print('>>>>>>>>>>>>>>> end evaluate case, epach=%d' % epoch)
def run_epoch(self, data, model, epoch, max_step=200):
start = time.time()
total_tokens = 0.
total_loss = 0.
tokens = 0.
# print('data=', len(data), data)
for i, batch in enumerate(data):
if i>=max_step:
break
out = model(batch.src, batch.trg, batch.src_mask, batch.trg_mask)
loss = self.compute_loss(out, batch.trg_y, batch.ntokens)
total_loss += loss
total_tokens += batch.ntokens
tokens += batch.ntokens
if i % 100 == 1:
elapsed = time.time() - start
print("Epoch %d Batch: %d Loss: %f Tokens per Sec: %fs" % (epoch, i - 1, loss / batch.ntokens, (tokens.float() / elapsed / 1000.)))
start = time.time()
tokens = 0
return total_loss / total_tokens
def evaluate(self, data):
"""
在data上用训练好的模型进行预测,打印模型翻译结果
"""
# 梯度清零
with torch.no_grad():
# 在data的英文数据长度上遍历下标
for _ in range(10):
i = random.randint(0, len(data.dev_en) - 1)
# print('i=', i)
# TODO: 打印待翻译的src句子
cn_sent = " ".join([data.cn_index_dict[w] for w in data.dev_cn[i]])
print("\n" + cn_sent)
# TODO: 打印对应的句子答案
en_sent = " ".join([data.en_index_dict[w] for w in data.dev_en[i]])
print("".join(en_sent))
# 将当前以单词id表示的英文句子数据转为tensor,并放如DEVICE中
src = torch.from_numpy(np.array(data.dev_cn[i])).long().to(DEVICE)
# 增加一维
src = src.unsqueeze(0)
# 设置attention mask
src_mask = (src != 0).unsqueeze(-2)
# 用训练好的模型进行decode预测
out = self.greedy_decode(src, src_mask, max_len=self.MAX_LENGTH, start_symbol=data.en_word_dict["BOS"])
# 初始化一个用于存放模型翻译结果句子单词的列表
translation = []
# 遍历翻译输出字符的下标(注意:开始符"BOS"的索引0不遍历)
for j in range(1, out.size(1)):
# 获取当前下标的输出字符
sym = data.en_index_dict[out[0, j].item()]
# 如果输出字符不为'EOS'终止符,则添加到当前句子的翻译结果列表
if sym != 'EOS':
translation.append(sym)
# 否则终止遍历
else:
break
# 打印模型翻译输出的中文句子结果
print("translation: %s" % " ".join(translation))
def greedy_decode(self, src, src_mask, max_len, start_symbol):
"""
传入一个训练好的模型,对指定数据进行预测
"""
# 先用encoder进行encode
memory = self.model.encode(src, src_mask)
# print('src=', src.shape, src)
# print('src_mask=', src_mask.shape, src_mask)
# print('memory=', memory.shape, memory)
# 初始化预测内容为1×1的tensor,填入开始符('BOS')的id,并将type设置为输入数据类型(LongTensor)
ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)
# print('ys=', ys.shape, ys)
# 遍历输出的长度下标
for i in range(max_len - 1):
# decode得到隐层表示
tgt_msk = subsequent_mask(ys.size(1)).type_as(src.data)
# print('i=',i,'tgt_msk=', tgt_msk)
out = self.model.decode(memory,
src_mask,
Variable(ys),
Variable(tgt_msk))
# 获取当前位置最大概率的预测词id
_, next_word = torch.max(out, dim=2)
# print('i = ', i, ', out=', out.shape, out)
# print('i = ', i, ', next_word=', next_word.shape, next_word)
next_word = next_word.data[0, 0]
# 将当前位置预测的字符id与之前的预测内容拼接起来
ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
# print('i = ', i, ' ,next_word=', next_word, ', ys=', ys)
return ys
if __name__=='__main__':
TRAIN_FILE = '../datas/eng-cmn-train.txt'
DEV_FILE = '../datas/eng-cmn-eval.txt'
TEST_FILE = ''
SAVE_FILE = 'model.pt'
h = 8
d_model = 512
d_ff = 2048
dropout = 0.1
N = 6
epochs = 52
batch_size = 64
print('<<<<<<<<<<<<<<<start one step>>>>>>>>>>>>>>>>>>>>>\n1.start prepare data++++++++++++++++++++++\n')
data = PrepareData(TRAIN_FILE, DEV_FILE, batch_size, DEVICE)
print("src_vocab %d" % data.src_vocab)
print("tgt_vocab %d" % data.tgt_vocab)
print('self.train_en.len = ', len(data.train_en))
print('self.dev_en.len = ', len(data.dev_en))
print('\n2.start to train++++++++++++++++++++++++\n')
transformer = Transformer(h, d_model, d_ff, dropout, N, data.src_vocab, data.tgt_vocab, DEVICE).to(DEVICE)
trainer = Trainer(epochs, data.tgt_vocab, d_model, transformer, SAVE_FILE='model.pt', MAX_LENGTH=20)
trainer.train(data)
print('\n3.start to eval++++++++++++++++++++++++\n')
transformer.load_state_dict(torch.load(SAVE_FILE))
# trainer = Trainer(epochs, data.tgt_vocab, d_model, transformer, SAVE_FILE='model1.pt', MAX_LENGTH=5)
trainer.evaluate(data)