Full Model
在这里,我们定义一个函数,它采用超参数并生成一个完整的模型。
def make_model(src_vocab, tgt_vocab, N=6,
d_model=512, d_ff=2048, h=8, dropout=0.1):
"提示:从超参数构建模型。"
c = copy.deepcopy
attn = MultiHeadedAttention(h, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
model = EncoderDecoder(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
Decoder(DecoderLayer(d_model, c(attn), c(attn),
c(ff), dropout), N),
nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
Generator(d_model, tgt_vocab))
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform(p)
return model
# Small example model.
tmp_model = make_model(10, 10, 2)
None
Training
本节介绍了我们模型的训练方法。
我们停下来快速介绍一下训练一个标准编码器解码器模型所需的一些工具。 首先,我们定义一个批处理对象,它包含用于训练的src和目标句子,以及构造掩码。
Batches and Masking
class Batch:
"在训练期间用掩码持有一批数据的对象。"
def __init__(self, src, trg=None, pad=0):
self.src = src
self.src_mask = (src != pad).unsqueeze(-2)
if trg is not None:
self.trg = trg[:, :-1]
self.trg_y = trg[:, 1:]
self.trg_mask = \
self.make_std_mask(self.trg, pad)
self.ntokens = (self.trg_y != pad).data.sum()
@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
接下来,我们创建一个通用的训练和评分功能来记录损失。 我们传入一个通用的损失计算函数,它也处理参数更新。
Training Loop
def run_epoch(data_iter, model, loss_compute):
"标准训练和记录函数"
start = time.time()
total_tokens = 0
total_loss = 0
tokens = 0
for i, batch in enumerate(data_iter):
out = model.forward(batch.src, batch.trg,
batch.src_mask, batch.trg_mask)
loss = loss_compute(out, batch.trg_y, batch.ntokens)
total_loss += loss
total_tokens += batch.ntokens
tokens += batch.ntokens
if i % 50 == 1:
elapsed = time.time() - start
print("Epoch Step: %d Loss: %f Tokens per Sec: %f" %
(i, loss / batch.ntokens, tokens / elapsed))
start = time.time()
tokens = 0
return total_loss / total_tokens
Training Data and Batching
我们使用标准的WMT 2014 English-German数据集进行了训练,该数据集包含大约450万个句子对。 句子使用字节对编码,该编码分享了大约37000个词的 source-target词汇表。 对于英语 - 法语,我们使用了大得多的 WMT 2014 English-French数据集,该数据集由3600万个句子组成,并将词分成32000个单词词汇。
句子对按照近似的序列长度进行批处理。 每个训练批包含一个句子对集合,其包含大约25000个来源词和25000个目标词。
我们将使用torch text进行批处理。 这将在下面更详细地讨论。 在这里,我们使用torchtext函数创建批次,确保填充到最大批量的大小不超过阈值(如果我们有8 gpus,则为25000)。
global max_src_in_batch, max_tgt_in_batch
def batch_size_fn(new, count, sofar):
"继续增加批量并计算词+填充的总数。"
global max_src_in_batch, max_tgt_in_batch
if count == 1:
max_src_in_batch = 0
max_tgt_in_batch = 0
max_src_in_batch = max(max_src_in_batch, len(new.src))
max_tgt_in_batch = max(max_tgt_in_batch, len(new.trg) + 2)
src_elements = count * max_src_in_batch
tgt_elements = count * max_tgt_in_batch
return max(src_elements, tgt_elements)
Hardware and Schedule
我们使用8块NVIDIA P100 GPU在一台机器上训练我们的模型。 对于本文所述的使用超参数的基本模型,每个训练步骤大约需要0.4秒。 我们对基础模型进行了总共100,000步或12小时的训练。 对于我们的大型模型,步进时间为1.0秒。 大型模型经过300,000步(3.5天)的培训。
Optimizer
我们采用Adam优化器,其中β1= 0.9,β2= 0.98,ε= 10-9。 我们根据公式改变了训练过程中的学习率:
这意味着在第一个预热训练步骤线性地增加学习速率,此后与步数的反平方根成比例地减小它。 我们使用了warmupsteps = 4000。
注意:这部分非常重要。 需要训练模型的这种设置。
class NoamOpt:
"执行速率的Optim包装器。"
def __init__(self, model_size, factor, warmup, optimizer):
self.optimizer = optimizer
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)))
def get_std_opt(model):
return NoamOpt(model.src_embed[0].d_model, 2, 4000,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
针对不同模型大小和优化超参数的模型的曲线示例。
# lrate超参数的三种设置。
opts = [NoamOpt(512, 1, 4000, None),
NoamOpt(512, 1, 8000, None),
NoamOpt(256, 1, 4000, None)]
plt.plot(np.arange(1, 20000), [[opt.rate(i) for opt in opts] for i in range(1, 20000)])
plt.legend(["512:4000", "512:8000", "256:4000"])
None
Regularization
Label Smoothing
在训练期间,我们采用ε= 0.1的标签平滑。 这会伤害困惑度,因为模型会更加不确定,但提高了准确性和BLEU分数。
我们使用KL div loss实现标签平滑。 我们不是使用one-hot目标分布,而是创建一个分布,它对正确的单词和其余整个词汇表中分布的平滑(smoothing)质量的置信度(confidence)。
class LabelSmoothing(nn.Module):
"实现标签平滑。"
def __init__(self, size, padding_idx, smoothing=0.0):
super(LabelSmoothing, self).__init__()
self.criterion = nn.KLDivLoss(size_average=False)
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
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))
在这里,我们可以看到一个基于置信度将质量分布到单词的示例。
# 标签平滑的示例。
crit = LabelSmoothing(5, 0, 0.4)
predict = torch.FloatTensor([[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0]])
v = crit(Variable(predict.log()),
Variable(torch.LongTensor([2, 1, 0])))
# Show the target distributions expected by the system.
plt.imshow(crit.true_dist)
None
crit = LabelSmoothing(5, 0, 0.1)
def loss(x):
d = x + 3 * 1
predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d],
])
#print(predict)
return crit(Variable(predict.log()),
Variable(torch.LongTensor([1]))).data[0]
plt.plot(np.arange(1, 100), [loss(x) for x in range(1, 100)])
None
A First Example
我们可以先尝试一个简单的复制任务。 给定来自小词汇表的随机输入符号集,目标是生成那些相同的符号。
Synthetic Data
def data_gen(V, batch, nbatches):
"为src-tgt复制任务生成随机数据。"
for i in range(nbatches):
data = torch.from_numpy(np.random.randint(1, V, size=(batch, 10)))
data[:, 0] = 1
src = Variable(data, requires_grad=False)
tgt = Variable(data, requires_grad=False)
yield Batch(src, tgt, 0)
Loss Computation
class SimpleLossCompute:
"简单的损失计算和训练功能。"
def __init__(self, generator, criterion, opt=None):
self.generator = generator
self.criterion = criterion
self.opt = opt
def __call__(self, x, y, norm):
x = self.generator(x)
loss = self.criterion(x.contiguous().view(-1, x.size(-1)),
y.contiguous().view(-1)) / norm
loss.backward()
if self.opt is not None:
self.opt.step()
self.opt.optimizer.zero_grad()
return loss.data[0] * norm
Greedy Decoding
# 训练简单的复制任务。
V = 11
criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
model = make_model(V, V, N=2)
model_opt = NoamOpt(model.src_embed[0].d_model, 1, 400,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
for epoch in range(10):
model.train()
run_epoch(data_gen(V, 30, 20), model,
SimpleLossCompute(model.generator, criterion, model_opt))
model.eval()
print(run_epoch(data_gen(V, 30, 5), model,
SimpleLossCompute(model.generator, criterion, None)))
为简单起见,此代码使用greedy decoding来预测转换。
def greedy_decode(model, src, src_mask, max_len, start_symbol):
memory = model.encode(src, src_mask)
ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)
for i in range(max_len-1):
out = model.decode(memory, src_mask,
Variable(ys),
Variable(subsequent_mask(ys.size(1))
.type_as(src.data)))
prob = model.generator(out[:, -1])
_, next_word = torch.max(prob, dim = 1)
next_word = next_word.data[0]
ys = torch.cat([ys,
torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
return ys
model.eval()
src = Variable(torch.LongTensor([[1,2,3,4,5,6,7,8,9,10]]) )
src_mask = Variable(torch.ones(1, 1, 10) )
print(greedy_decode(model, src, src_mask, max_len=10, start_symbol=1))
A Real World Example
现在我们考虑使用IWSLT德语 - 英语翻译任务的真实示例。 此任务比文章中考虑的WMT任务小得多,但它展示了整个系统。 我们还展示了如何使用multi-gpu处理使计算非常快。
#!pip install torchtext spacy
#!python -m spacy download en
#!python -m spacy download de
Data Loading
我们将使用torchtext和spacy对加载的数据集进行标记化。
# 用于数据加载。
from torchtext import data, datasets
if True:
import spacy
spacy_de = spacy.load('de')
spacy_en = spacy.load('en')
def tokenize_de(text):
return [tok.text for tok in spacy_de.tokenizer(text)]
def tokenize_en(text):
return [tok.text for tok in spacy_en.tokenizer(text)]
BOS_WORD = ''
EOS_WORD = ''
BLANK_WORD = ""
SRC = data.Field(tokenize=tokenize_de, pad_token=BLANK_WORD)
TGT = data.Field(tokenize=tokenize_en, init_token = BOS_WORD,
eos_token = EOS_WORD, pad_token=BLANK_WORD)
MAX_LEN = 100
train, val, test = datasets.IWSLT.splits(
exts=('.de', '.en'), fields=(SRC, TGT),
filter_pred=lambda x: len(vars(x)['src']) <= MAX_LEN and
len(vars(x)['trg']) <= MAX_LEN)
MIN_FREQ = 2
SRC.build_vocab(train.src, min_freq=MIN_FREQ)
TGT.build_vocab(train.trg, min_freq=MIN_FREQ)
批量处理对速度很重要。 我们希望拥有非常均匀的分批,绝对最小的填充。 要做到这一点,我们必须破解默认的torchtext批处理。 此代码修补其默认批处理,以确保我们搜索足够的句子来查找紧密的批处理。
Iterators
class MyIterator(data.Iterator):
def create_batches(self):
if self.train:
def pool(d, random_shuffler):
for p in data.batch(d, self.batch_size * 100):
p_batch = data.batch(
sorted(p, key=self.sort_key),
self.batch_size, self.batch_size_fn)
for b in random_shuffler(list(p_batch)):
yield b
self.batches = pool(self.data(), self.random_shuffler)
else:
self.batches = []
for b in data.batch(self.data(), self.batch_size,
self.batch_size_fn):
self.batches.append(sorted(b, key=self.sort_key))
def rebatch(pad_idx, batch):
"Fix order in torchtext to match ours"
src, trg = batch.src.transpose(0, 1), batch.trg.transpose(0, 1)
return Batch(src, trg, pad_idx)
Multi-GPU Training
最后目标要真正快速训练,我们将使用multi-gpu。 此代码实现了多gpu字生成。 它不是Transformer特有的,所以不会详细介绍。 这个想法是将训练时的单词生成分成块,以便在许多不同的gpus上并行处理。 我们使用pytorch并行原语来做到这一点:
replicate - 将模块拆分到不同的gpus上。
scatter - 将批次分到不同的gpus
parallel_apply - 将模块应用于不同gpus上的批处理
gather - 将分散的数据聚回到一个gpu上。
nn.DataParallel - 一个特殊的模块包装器,在评估之前调用它们。
# Skip if not interested in multigpu.
class MultiGPULossCompute:
"A multi-gpu loss compute and train function."
def __init__(self, generator, criterion, devices, opt=None, chunk_size=5):
# Send out to different gpus.
self.generator = generator
self.criterion = nn.parallel.replicate(criterion,
devices=devices)
self.opt = opt
self.devices = devices
self.chunk_size = chunk_size
def __call__(self, out, targets, normalize):
total = 0.0
generator = nn.parallel.replicate(self.generator,
devices=self.devices)
out_scatter = nn.parallel.scatter(out,
target_gpus=self.devices)
out_grad = [[] for _ in out_scatter]
targets = nn.parallel.scatter(targets,
target_gpus=self.devices)
# Divide generating into chunks.
chunk_size = self.chunk_size
for i in range(0, out_scatter[0].size(1), chunk_size):
# Predict distributions
out_column = [[Variable(o[:, i:i+chunk_size].data,
requires_grad=self.opt is not None)]
for o in out_scatter]
gen = nn.parallel.parallel_apply(generator, out_column)
# Compute loss.
y = [(g.contiguous().view(-1, g.size(-1)),
t[:, i:i+chunk_size].contiguous().view(-1))
for g, t in zip(gen, targets)]
loss = nn.parallel.parallel_apply(self.criterion, y)
# Sum and normalize loss
l = nn.parallel.gather(loss,
target_device=self.devices[0])
l = l.sum()[0] / normalize
total += l.data[0]
# Backprop loss to output of transformer
if self.opt is not None:
l.backward()
for j, l in enumerate(loss):
out_grad[j].append(out_column[j][0].grad.data.clone())
# Backprop all loss through transformer.
if self.opt is not None:
out_grad = [Variable(torch.cat(og, dim=1)) for og in out_grad]
o1 = out
o2 = nn.parallel.gather(out_grad,
target_device=self.devices[0])
o1.backward(gradient=o2)
self.opt.step()
self.opt.optimizer.zero_grad()
return total * normalize
现在我们建立了 model, criterion, optimizer, data iterators, and paralelization。
# GPUs to use
devices = [0, 1, 2, 3]
if True:
pad_idx = TGT.vocab.stoi[""]
model = make_model(len(SRC.vocab), len(TGT.vocab), N=6)
model.cuda()
criterion = LabelSmoothing(size=len(TGT.vocab), padding_idx=pad_idx, smoothing=0.1)
criterion.cuda()
BATCH_SIZE = 12000
train_iter = MyIterator(train, batch_size=BATCH_SIZE, device=0,
repeat=False, sort_key=lambda x: (len(x.src), len(x.trg)),
batch_size_fn=batch_size_fn, train=True)
valid_iter = MyIterator(val, batch_size=BATCH_SIZE, device=0,
repeat=False, sort_key=lambda x: (len(x.src), len(x.trg)),
batch_size_fn=batch_size_fn, train=False)
model_par = nn.DataParallel(model, device_ids=devices)
None
现在我们训练模型。 我会稍微使用一些预热步骤,但其他一切都使用默认参数。 在具有4个特斯拉V100的AWS p3.8xlarge上,每秒运行约27,000个词,批量大小为12,000。
Training the System
#!wget https://s3.amazonaws.com/opennmt-models/iwslt.pt
if False:
model_opt = NoamOpt(model.src_embed[0].d_model, 1, 2000,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
for epoch in range(10):
model_par.train()
run_epoch((rebatch(pad_idx, b) for b in train_iter),
model_par,
MultiGPULossCompute(model.generator, criterion,
devices=devices, opt=model_opt))
model_par.eval()
loss = run_epoch((rebatch(pad_idx, b) for b in valid_iter),
model_par,
MultiGPULossCompute(model.generator, criterion,
devices=devices, opt=None))
print(loss)
else:
model = torch.load("iwslt.pt")
经过训练,我们可以解码模型以生成一组翻译。 在这里,我们只需翻译验证集中的第一句话。 此数据集非常小,因此使用贪婪搜索的翻译相当准确。
for i, batch in enumerate(valid_iter):
src = batch.src.transpose(0, 1)[:1]
src_mask = (src != SRC.vocab.stoi[""]).unsqueeze(-2)
out = greedy_decode(model, src, src_mask,
max_len=60, start_symbol=TGT.vocab.stoi[""])
print("Translation:", end="\t")
for i in range(1, out.size(1)):
sym = TGT.vocab.itos[out[0, i]]
if sym == "": break
print(sym, end =" ")
print()
print("Target:", end="\t")
for i in range(1, batch.trg.size(0)):
sym = TGT.vocab.itos[batch.trg.data[i, 0]]
if sym == "": break
print(sym, end =" ")
print()
break
Translation: . In my language , that means , thank you very much .
Gold: . It means in my language , thank you very much .
Additional Components: BPE, Search, Averaging
所以这主要涵盖了Transformer模型本身。 我们没有明确涵盖四个方面。 我们还在OpenNMT-py中实现了所有这些附加功能。
1)BPE / Word-piece:我们可以使用库首先将数据预处理为子字单元。 参见Rico Sennrich的subword-nmt实现。 这些模型将训练数据转换为如下所示:
▁Die ▁Protokoll datei ▁kann ▁ heimlich ▁per ▁E - Mail ▁oder ▁FTP ▁an ▁einen ▁bestimmte n ▁Empfänger ▁gesendet ▁werden .
2)Shared Embeddings:当使用具有共享词汇表的BPE时,我们可以在源/目标/生成器之间共享相同的权重向量。 有关详细信息,请参阅(https://arxiv.org/abs/1608.05859)。 要将其添加到模型,只需执行以下操作:
if False:
model.src_embed[0].lut.weight = model.tgt_embeddings[0].lut.weight
model.generator.lut.weight = model.tgt_embed[0].lut.weight
4)Model Averaging:本文平均最后k个检查点以创建一个集合效果。 如果我们有一堆模型,我们可以这样做:
def average(model, models):
"Average models into model"
for ps in zip(*[m.params() for m in [model] + models]):
p[0].copy_(torch.sum(*ps[1:]) / len(ps[1:]))
Results
在WMT 2014English-to-German翻译任务中,大Transformer模型(表2中的Transformer(大))优于先前报告的最佳模型(包ensembles)超过2.0 BLEU,建立了一个新的最佳BLEU得分28.4。 该模型的配置列于表3的底部.在8个P100 GPU上的训练耗时3.5天。 甚至我们的基础模型也超过了之前发布的所有模型,而且训练成本是其他有竞争力的模型的一小部分。
在WMT 2014English-to-French翻译任务中,我们的大型模型获得了41.0的BLEU分数,优于以前发布的所有单一模型,不到以前最先进模型训练成本的1/4 。 针对英语到法语训练的Transformer(大)模型使用Pdrop = 0.1,而不是0.3。
我们在这里编写的代码是基本模型的一个版本。 完整训练版本的模型如下(http://opennmt.net/Models-py/)。
通过最后一节中的扩展,OpenNMT-py在EN-DE WMT上达到26.9。 在这里,我已将这些参数加载到我们的重现中。
!wget https://s3.amazonaws.com/opennmt-models/en-de-model.pt
model, SRC, TGT = torch.load("en-de-model.pt")
model.eval()
sent = "▁The ▁log ▁file ▁can ▁be ▁sent ▁secret ly ▁with ▁email ▁or ▁FTP ▁to ▁a ▁specified ▁receiver".split()
src = torch.LongTensor([[SRC.stoi[w] for w in sent]])
src = Variable(src)
src_mask = (src != SRC.stoi[""]).unsqueeze(-2)
out = greedy_decode(model, src, src_mask,
max_len=60, start_symbol=TGT.stoi[""])
print("Translation:", end="\t")
trans = " "
for i in range(1, out.size(1)):
sym = TGT.itos[out[0, i]]
if sym == "": break
trans += sym + " "
print(trans)
Translation: ▁Die ▁Protokoll datei ▁kann ▁ heimlich ▁per ▁E - Mail ▁oder ▁FTP ▁an ▁einen ▁bestimmte n ▁Empfänger ▁gesendet ▁werden .
Attention Visualization
即使使用greedy解码器,翻译的也不错。 我们可以进一步想象它,看看每个注意力层发生了什么。
tgt_sent = trans.split()
def draw(data, x, y, ax):
seaborn.heatmap(data,
xticklabels=x, square=True, yticklabels=y, vmin=0.0, vmax=1.0,
cbar=False, ax=ax)
for layer in range(1, 6, 2):
fig, axs = plt.subplots(1,4, figsize=(20, 10))
print("Encoder Layer", layer+1)
for h in range(4):
draw(model.encoder.layers[layer].self_attn.attn[0, h].data,
sent, sent if h ==0 else [], ax=axs[h])
plt.show()
for layer in range(1, 6, 2):
fig, axs = plt.subplots(1,4, figsize=(20, 10))
print("Decoder Self Layer", layer+1)
for h in range(4):
draw(model.decoder.layers[layer].self_attn.attn[0, h].data[:len(tgt_sent), :len(tgt_sent)],
tgt_sent, tgt_sent if h ==0 else [], ax=axs[h])
plt.show()
print("Decoder Src Layer", layer+1)
fig, axs = plt.subplots(1,4, figsize=(20, 10))
for h in range(4):
draw(model.decoder.layers[layer].self_attn.attn[0, h].data[:len(tgt_sent), :len(sent)],
sent, tgt_sent if h ==0 else [], ax=axs[h])
plt.show()
Conclusion
希望这段代码对未来的研究很有用。 如果您有任何问题,请与我们联系。 如果您发现此代码有用,请查看我们的其他OpenNMT工具。
@inproceedings{opennmt,
author = {Guillaume Klein and
Yoon Kim and
Yuntian Deng and
Jean Senellart and
Alexander M. Rush},
title = {OpenNMT: Open-Source Toolkit for Neural Machine Translation},
booktitle = {Proc. ACL},
year = {2017},
url = {https://doi.org/10.18653/v1/P17-4012},
doi = {10.18653/v1/P17-4012}
}
Cheers, srush