Keras微调BERT模型时无早停配置却疑似跳过最后一个epoch

实验配置说明 #

• 训练总轮次设置为2，batch size为24，未配置任何早停类回调函数
• 训练基于TensorFlow Hub加载的BERT编码器搭建微调流程，采用AdamW优化器，初始学习率2e-5
• 每轮训练的步数计算方式为train_df.shape[0] // batch_size
• 总训练步数为单轮步数乘以总轮次，warmup步数为总训练步数的1/10

异常现象 #

第一个epoch完整跑完，单步耗时652ms，总耗时8380s；但第二个epoch显示所有12816步执行完成时，总耗时仅为0微秒，loss值没有对应实际训练的计算过程，疑似直接跳过了整轮训练。提问方核对过全流程未配置任何早停相关逻辑，怀疑是否优化器配置存在错误。

核心实现代码 #

epochs = 2
batch_size = 24
init_lr = 2e-5
print(f'Fine tuning {tfhub_handle_encoder} model')
steps_per_epoch = train_df.shape[0] // batch_size
num_train_steps = steps_per_epoch * epochs
num_warmup_steps = num_train_steps // 10
validation_steps = val_df.shape[0] // batch_size
AUTOTUNE = tf.data.AUTOTUNE

train_set = tf.data.Dataset.from_tensor_slices(((train_df['query'].values,train_df['doc'].values),np.array(train_labels)))
train_set = train_set.batch(24)
train_set = train_set.cache().prefetch(buffer_size=AUTOTUNE)

preprocess_inputs = preprocessor.inputs   # BERT预处理器的输入
bert_input_data = preprocessor(preprocess_inputs)  # 预处理器输出
bert_output = bert(bert_input_data)  # 预处理器输出接入BERT模型
encoder_output = encoder(bert_output) # 获取BERT编码输出
reconstruction = decoder(encoder_output) # 编码输出接入解码器
bert_qe_vae = tf.keras.Model(preprocess_inputs, [encoder_output, reconstruction])
all_losses = {'encoder_output': kl_divergence, 'reconstruction': recon_loss}
optimizer = optimization.create_optimizer(
  init_lr=init_lr,
  num_train_steps=num_train_steps,
  num_warmup_steps=num_warmup_steps,
  optimizer_type='adamw')
bert_qe_vae.compile(optimizer=optimizer, loss=all_losses)
bert_qe_vae.fit(x=train_set,
        epochs=epochs,
        steps_per_epoch=steps_per_epoch)

训练日志输出 #

正在微调模型：bert_en_uncased_L-12_H-768_A-12/3
Epoch 1/2
12816/12816 [==============================] - 8380s 652ms/step - loss: 647.3066 - encoder_loss: 83.6949 - decoder_loss: 563.6116 
Epoch 2/2
12816/12816 [==============================] - 0s 2us/step - loss: 549.6426 - encoder_loss: 19.1478 - decoder_loss: 530.4949 

内容的提问来源于stack exchange，提问作者Wei Ou