前言
还在为机器翻译模型从理论到落地卡壳?系列博客第三弹——模型训练篇强势登场,手把手带你走完Transformer中日翻译项目的最后关键一步!
前两期我们搞定了数据预处理(分词、词表构建全流程)和模型搭建(词嵌入、位置编码、编码器解码器核心结构),而这一篇,将聚焦让模型“学会翻译”的核心秘籍:
- 如何设计损失函数,让模型精准捕捉中日语言差异?
- 优化器参数怎么调,才能让训练更稳定、收敛更快?
从数据到模型,再到训练全流程,一套代码跑通Transformer翻译任务。
本篇主要介绍机器翻译项目中模型训练部分是如何处理的,延续之前篇章预处理篇,模型篇的内容。
系列博客第一弹——基于Transformer的机器翻译——预处理篇
系列博客第二弹——基于Transformer的机器翻译——模型篇
系列博客第三弹——基于Transformer的机器翻译——训练篇
1.模型介绍
1.1模型参数配置
参数主要有源语言与目标语言词表大小
、词嵌入维度
、多头注意力头数
、前馈神经网络隐藏层维度
、编码器与解码器层数
、批量大小
等,这些均属于超参数,需要在训练中不断调整,考虑到自身硬件条件,超参数设置如下:
# 模型参数配置
SRC_VOCAB_SIZE = len(ja_vocab) # 源语言(日语)词汇表大小
TGT_VOCAB_SIZE = len(ch_vocab) # 目标语言(中文)词汇表大小
EMB_SIZE = 512 # 词嵌入维度(与Transformer的d_model一致)
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
)
1.2模型结构
可以直接打印模型,查看模型结构,如下:
transformer
输出结果:
Seq2SeqTransformer((transformer_encoder): TransformerEncoder((layers): ModuleList((0-2): 3 x TransformerEncoderLayer((self_attn): MultiheadAttention((out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True))(linear1): Linear(in_features=512, out_features=512, bias=True)(dropout): Dropout(p=0.1, inplace=False)(linear2): Linear(in_features=512, out_features=512, bias=True)(norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(dropout1): Dropout(p=0.1, inplace=False)(dropout2): Dropout(p=0.1, inplace=False))))(transformer_decoder): TransformerDecoder((layers): ModuleList((0-2): 3 x TransformerDecoderLayer((self_attn): MultiheadAttention((out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True))(multihead_attn): MultiheadAttention((out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True))(linear1): Linear(in_features=512, out_features=512, bias=True)(dropout): Dropout(p=0.1, inplace=False)(linear2): Linear(in_features=512, out_features=512, bias=True)(norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(norm3): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(dropout1): Dropout(p=0.1, inplace=False)(dropout2): Dropout(p=0.1, inplace=False)(dropout3): Dropout(p=0.1, inplace=False))))(generator): Linear(in_features=512, out_features=26854, bias=True)(src_tok_emb): TokenEmbedding((embedding): Embedding(24058, 512))(tgt_tok_emb): TokenEmbedding((embedding): Embedding(26854, 512))(positional_encoding): PositionalEncoding((dropout): Dropout(p=0.1, inplace=False))
)
1.3参数初始化
一般主要有两种初始化模型参数的方法,包括He初始化(即一般默认的初始化方法)
和Xavier初始化方法
。
- He初始化:适用于使用
ReLU
或其变种Leaky ReLU
这类激活函数的神经网络。 - Xavier初始化:适用于
tanh
、sigmoid
等这类对称激活函数。
这里采用Xavier初始化
,代码如下:
# 参数初始化(Xavier均匀初始化,缓解梯度消失/爆炸)
for p in transformer.parameters():if p.dim() > 1: # 仅初始化非标量参数nn.init.xavier_uniform_(p)
为了加快模型的训练速度,这里使用GPU加速,即将模型迁移至GPU上进行训练。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transformer = transformer.to(device)
如果未配置GPU环境,此处会默认使用CPU。
2.训练策略
2.1损失函数
因为该问题最终是一个分类问题,预测词表中的一个单词。而分类问题,一般采用交叉熵损失函数。同时,因为此前我们填充<PAD_IDX>
字符使输入句子长度保持一致,所以在计算交叉熵的时候,应该将该字符产生的损失忽略。
代码如下:
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)
2.2优化器
优化器大都会选用自适应优化器Adam,相对于常规的SGD(随机梯度下降方法),避免了频繁调整参数,同时能够更快的收敛。代码如下:
optimizer = torch.optim.Adam(transformer.parameters(), # 待优化的模型参数lr=0.0001, # 学习率(控制参数更新步长)betas=(0.9, 0.98), # 动量参数(控制历史梯度的衰减)eps=1e-9 # 数值稳定性参数(防止除零)
)
2.3训练过程
整体的训练过程主要包括五个步骤:
- 梯度清零
- 计算模型输出
- 计算损失
- 方向传播
- 参数更新
模型输入包括两部分,输入的源语言,和目标语言已经生成的部分,分别对应于编码器的输入和解码器的输入。
因为是逐个token
进行预测,所以编码器的输入不包括句子的最后一个token
即EOS_IDX
,模型预测输出到EOS_IDX
即结束,相应的监督标签为目标语言不包括第一个token即BOS_IDX
。
相应代码为:
src = src.to(device) # 将源序列移动到GPU/CPUtgt = tgt.to(device) # 将目标序列移动到GPU/CPUtgt_input = tgt[:-1, :] # 目标输入去掉最后一个词tgt_output = tgt[1:, :] # 目标输出去掉第一个词
因为模型在进行训练时,不能关注到目标输入后续的token
,因此需要掩码注意力机制
2.3.1顺序掩码
采用下三角掩码,防止模型在预测第i个词时关注到第i+1、i+2… 个词(即 “未来信息”)
- 通过
torch.triu(torch.ones((3,3)))
生成上三角矩阵(对角线及以上为 1,其余为 0) ==1
转为布尔矩阵(True 表示原位置为 1),再transpose(0,1)
转置(行变列,列变行)masked_fill
填充:True 位置(有效)填 0.0,False 位置(无效,未来信息)填-inf
本案例主要通过函数generate_square_subsequent_mask
函数实现,代码如下:
def generate_square_subsequent_mask(sz):# 1. 生成上三角矩阵(对角线及以上为1,其余为0)mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)# 2. 将0的位置填充为-inf(无效位置),1的位置填充为0(有效位置)mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))return mask
2.3.2完整掩码组合
这里主要4 种掩码,分别用于源序列和目标序列的注意力计算。
2.3.2.1源序列自注意力掩码
源序列是输入,无需掩盖 “未来信息”,因此全为 0(允许关注所有位置)。
2.3.2.2目标序列自注意力掩码
可以通过调用 generate_square_subsequent_mask
实现。
2.3.2.3源序列填充掩码
标记源序列中PAD
的位置,让模型在注意力计算时忽略这些无效位置。
2.3.2.4目标序列填充掩码
类似源序列,标记目标序列中PAD的位置。
2.3.2.5完整代码
上述四种掩码的生成代码可表述为以下代码:
def create_mask(src, tgt):src_seq_len = src.shape[0] # 源序列长度(seq_len)tgt_seq_len = tgt.shape[0] # 目标序列长度(seq_len)# 目标序列的注意力掩码(下三角掩码,防止关注未来词)tgt_mask = generate_square_subsequent_mask(tgt_seq_len)# 源序列的注意力掩码(全0,允许关注所有位置)src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)# 源序列的填充掩码(标记<pad>的位置)src_padding_mask = (src == PAD_IDX).transpose(0, 1)# 目标序列的填充掩码(标记<pad>的位置)tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask
2.3.3训练代码
由上述代码可得,训练过程可表述为以下代码:
- 获取当前批量的输入与目标输出序列
- 生成掩码
- 计算模型预测输出
- 清空梯度
- 计算预测输出与目标输出之间的损失
- 方向传播
- 更新梯度
for idx, (src, tgt) in enumerate(train_iter): # 遍历训练数据迭代器src = src.to(device) # 将源序列移动到GPU/CPUtgt = tgt.to(device) # 将目标序列移动到GPU/CPUtgt_input = tgt[:-1, :] # 目标输入去掉最后一个词tgt_output = 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]), # 展平预测结果([seq_len*batch_size, vocab_size])tgt_out.reshape(-1) # 展平真实标签([seq_len*batch_size]))loss.backward() # 反向传播计算梯度optimizer.step() # 更新模型参数losses += loss.item() # 累加单批损失
为了减小梯度抖动,采用当前iter
内所有损失的平均和,同时为了可视化损失折线图,加入了tensorboard
可视化,并加入进度条来显示训练过程,当前训练过程。代码如下:
def train_epoch(model, train_iter, optimizer):model.train() # 开启训练模式(启用Dropout等)losses = 0 # 累计损失值for idx, (src, tgt) in enumerate(train_iter): # 遍历训练数据迭代器src = src.to(device) # 将源序列移动到GPU/CPUtgt = tgt.to(device) # 将目标序列移动到GPU/CPUtgt_input = tgt[:-1, :] # 目标输入去掉最后一个词tgt_output = 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]), # 展平预测结果([seq_len*batch_size, vocab_size])tgt_out.reshape(-1) # 展平真实标签([seq_len*batch_size]))# 写入tensorboardwriter.add_scalar('train_loss',loss.item(),idx)loss.backward() # 反向传播计算梯度optimizer.step() # 更新模型参数losses += loss.item() # 累加单批损失train_bar.set_postfix({'loss': '{:.6f}'.format(loss.item())})train_bar.update()return losses / len(train_iter) # 返回平均训练损失
为了避免长时间等待,因此使用当前数据集的10%,其中9%作为训练集,1%作为测试集。代码如下:
train_iter = DataLoader(train_data[:int(len(train_data)*0.09)], batch_size=BATCH_SIZE,shuffle=True, collate_fn=generate_batch)
test_iter = DataLoader(train_data[int(len(train_data)*0.09):int(len(train_data)*0.1)], batch_size=BATCH_SIZE,shuffle=True, collate_fn=generate_batch)
3.训练结果
假如此时按照每个iter
来进行记录损失(上述代码就是如此),结果会出现剧烈波动,整体还是会保持下降趋势,结果如下:
此时将其改成每个epoch来记录损失,取当前epoch的平均值,结果如下:
此时只统计了12个点,当然你可以通过增加epoch
轮数来增加统计的点数,但是曲线的变化趋势往往是判断模型过拟合和欠拟合的关键,因此一般采用每多个iter
内的加权值来统计一次损失。这里为了方便起见,采用的仍是按照epoch
进行统计的。
同时为了可视化训练进度,此处加入了进度条来实时展示训练的进度。统计损失使用tensorboard
进行可视化。
完整训练过程代码为:
训练函数:
writer=SummaryWriter('./logs')
def train_epoch(model, train_iter, optimizer,epoch):model.train() # 开启训练模式(启用Dropout等)losses = 0 # 累计损失值for idx, (src, tgt) in enumerate(train_iter): # 遍历训练数据迭代器src = src.to(device) # 将源序列移动到GPU/CPUtgt = tgt.to(device) # 将目标序列移动到GPU/CPUtgt_input = tgt[:-1, :] # 目标输入去掉最后一个词tgt_output = 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]), # 展平预测结果([seq_len*batch_size, vocab_size])tgt_out.reshape(-1) # 展平真实标签([seq_len*batch_size]))loss.backward() # 反向传播计算梯度optimizer.step() # 更新模型参数losses += loss.item() # 累加单批损失train_bar.set_postfix({'loss': '{:.6f}'.format(loss.item())})train_bar.update()# 写入tensorboardwriter.add_scalar('train_loss',losses / len(train_iter),epoch)return losses / len(train_iter) # 返回平均训练损失
验证函数:
def evaluate(model, test_iter,epoch):model.eval()losses = 0for src, tgt in test_iter:src = src.to(device)tgt = tgt.to(device)tgt_input = tgt[:-1, :]tgt_output = 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()test_bar.set_postfix({'loss': '{:.6f}'.format(loss.item())})test_bar.update()writer.add_scalar('test_loss',losses / len(test_iter),epoch)return losses / len(test_iter)
调用函数:
from tqdm.notebook import tqdm
epoch_bar = tqdm(desc='training routine',total=NUM_EPOCHS,position=0)train_bar = tqdm(desc='split=train',total=len(train_iter),position=0,leave=True)
test_bar=tqdm(desc='split=test',total=len(test_iter),position=0,leave=True)
best_loss = float('inf')
for epoch in range(1, NUM_EPOCHS+1):start_time = time.time()train_loss = train_epoch(transformer, train_iter, optimizer,epoch)test_loss = evaluate(transformer, test_iter,epoch)if test_loss < best_loss:best_loss = test_loss# 保存时分开保存torch.save(transformer.state_dict(), './share/model_weights.pth') # 只保存权重torch.save({'ja_vocab': ja_vocab, 'ch_vocab': ch_vocab}, './share/vocab.pth') # 保存词表print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Model saved.")epoch_bar.set_postfix({'train_loss': '{:.6f}'.format(train_loss)})epoch_bar.update()train_bar.n = 0test_bar.n=0end_time = time.time()print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "f"Epoch time = {(end_time - start_time):.3f}s"))
训练结果:
结语
至此,基于Transformer的机器翻译任务就介绍完毕了,至于模型性能的评价指标如BLEU并未做介绍,感兴趣的可以自行探索,希望本案例,能够对你有所帮助,感谢支持!
系列博客第一弹——基于Transformer的机器翻译——预处理篇
系列博客第二弹——基于Transformer的机器翻译——模型篇
系列博客第三弹——基于Transformer的机器翻译——训练篇