训练Transformer模型

236 阅读5分钟

样本数据准备

为了验证,这次创建一个虚拟数据集。然而,在实际情况下,将使用更大的数据集,并且该过程将涉及文本预处理以及为源语言和目标语言创建词汇映射。

  1. src_vocab_size, tgt_vocab_size: 源和目标序列的词汇表大小,均设置为5000。
  2. d_model: 模型嵌入的维度,设置为512。
  3. num_heads: 多头注意力机制中的注意力头数,设置为8。
  4. num_layers: 编码器和解码器的层数,设置为6。
  5. d_ff: 前馈网络中内层的维度,设置为2048。
  6. max_seq_length: 位置编码的最大序列长度,设置为100。
  7. dropout: 用于正则化的dropout率,设置为0.1。
src_vocab_size = 5000  
tgt_vocab_size = 5000  
d_model = 512  
num_heads = 8  
num_layers = 6  
d_ff = 2048  
max_seq_length = 100  
dropout = 0.1  
  
# 生成随机样本数据  
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)  
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)  

  1. src_data: 介于1和src_vocab_size之间的随机整数,表示一批源序列,形状为(64, max_seq_length)。
  2. tgt_data: 介于1和tgt_vocab_size之间的随机整数,表示一批目标序列,形状为(64, max_seq_length)。
  3. 这些随机序列可以用作输入到Transformer模型中,模拟了一批具有64个示例和长度为100的序列的数据。

创建Transformer实例

创建一个Transformer类的实例,用给定的超参数初始化它。实例将具有由这些超参数定义的架构和行为。

transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)  

初始化一个Transformer模型并生成可以输入到模型中的随机源和目标序列。选择的超参数决定了Transformer的具体结构和属性。这种设置可以是更大脚本的一部分,其中模型在实际的序列到序列任务上进行训练和评估,例如机器翻译或文本摘要。

训练模型

使用样本数据训练模型。在现实场景中,将使用一个更大的数据集,通常会被划分为训练和验证的独立集合。

损失函数和优化器

  1. criterion = nn.CrossEntropyLoss(ignore_index=0): 定义损失函数为交叉熵损失。ignore_index参数设置为0,意味着损失将不考虑索引为0的目标(通常保留给填充标记)。
  2. optimizer = optim.Adam(...): 定义优化器为Adam,学习率为0.0001和特定的beta值。

训练模式

transformer.train(): 将Transformer模型设置为训练模式,启用仅在训练期间应用的行为,如dropout。

训练循环模型100个周期

  1. for epoch in range(100): 遍历100个训练周期。
  2. optimizer.zero_grad(): 清除上一次迭代的梯度。
  3. output = transformer(src_data, tgt_data[:, :-1]): 将源数据和目标数据(每个序列中除去最后一个标记)传递通过Transformer。这在序列到序列任务中很常见,目标会逐个标记地移动。
  4. loss = criterion(...): 计算模型预测和目标数据(每个序列中除去第一个标记)之间的损失。损失通过将数据重塑为一维张量并使用交叉熵损失函数来计算。
  5. loss.backward(): 计算损失相对于模型参数的梯度。
  6. optimizer.step(): 使用计算出的梯度更新模型的参数。
  7. print(f"Epoch: {epoch+1}, Loss: {loss.item()}"): 打印当前周期号和该周期的损失值。
criterion = nn.CrossEntropyLoss(ignore_index=0)  
optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)  
  
transformer.train()  
  
for epoch in range(100):  
    optimizer.zero_grad()  
    output = transformer(src_data, tgt_data[:, :-1])  
    loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:].contiguous().view(-1))  
    loss.backward()  
    optimizer.step()  
    print(f"Epoch: {epoch+1}, Loss: {loss.item()}")  

在随机生成的源和目标序列上训练Transformer模型100个周期。它使用Adam优化器和交叉熵损失函数。每个周期的损失被打印出来,允许你监控训练进度。在现实场景中,用你的任务中的实际数据替换随机源和目标序列,例如机器翻译。

评估模型

在随机生成的验证数据集上评估Transformer模型,计算验证损失,并打印它。在现实场景中,随机验证数据应替换为你正在处理的任务的实际验证数据。验证损失可以给你一个模型在未见过的数据上表现如何的指示,这是衡量模型泛化能力的关键指标。

评估模式:

transformer.eval(): 将Transformer模型置于评估模式。这很重要,因为它关闭了仅在训练期间使用的某些行为,如dropout。

生成随机验证数据:

  1. val_src_data: 介于1和src_vocab_size之间的随机整数,表示一批验证源序列,形状为(64, max_seq_length)。
  2. val_tgt_data: 介于1和tgt_vocab_size之间的随机整数,表示一批验证目标序列,形状为(64, max_seq_length)。

验证循环:

  1. with torch.no_grad(): 禁用梯度计算,因为在验证期间我们不需要计算梯度。这可以减少内存消耗并加速计算。
  2. val_output = transformer(val_src_data, val_tgt_data[:, :-1]): 将验证源数据和验证目标数据(每个序列中除去最后一个标记)传递通过Transformer。
  3. val_loss = criterion(...): 计算模型预测和验证目标数据(每个序列中除去第一个标记)之间的损失。损失通过将数据重塑为一维张量并使用之前定义的交叉熵损失函数来计算。
  4. print(f"Validation Loss: {val_loss.item()}"): 打印验证损失值。
transformer.eval()  
  
# 生成随机样本验证数据  
val_src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)  
val_tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)  
  
with torch.no_grad():  
    val_output = transformer(val_src_data, val_tgt_data[:, :-1])  
    val_loss = criterion(val_output.contiguous().view(-1, tgt_vocab_size), val_tgt_data[:, 1:].contiguous().view(-1))  
    print(f"Validation Loss: {val_loss.item()}")