新手求助：无model.fit()时如何用train_step保存最佳模型或适配model.fit()？

我是机器学习新手，正在训练Transformer图像Caption模型，想在30个epochs里只保存性能最好的模型权重，但现在的代码会保存全部30个模型：

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
    name='train_accuracy')

transformer = Transformer(num_layer, d_model, num_heads, dff, row_size, col_size, target_vocab_size,
                          max_pos_encoding=target_vocab_size, rate=dropout_rate)

@tf.function
def train_step(img_tensor, tar):
    tar_inp = tar[:, :-1]
    tar_real = tar[:, 1:]

    dec_mask = create_masks_decoder(tar_inp)

    with tf.GradientTape() as tape:
        predictions, _ = transformer(img_tensor, tar_inp,
                                     True,
                                     dec_mask)
        loss = loss_function(tar_real, predictions)

    gradients = tape.gradient(loss, transformer.trainable_variables)
    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))

    train_loss(loss)
    train_accuracy(tar_real, predictions)

for epoch in range(30):
    start = time.time()

    train_loss.reset_states()
    train_accuracy.reset_states()

    for (batch, (img_tensor, tar)) in enumerate(dataset):
        train_step(img_tensor, tar)

        if batch % 50 == 0:
            print('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(
                epoch + 1, batch, train_loss.result(), train_accuracy.result()))

    print('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1,
                                                        train_loss.result(),
                                                        train_accuracy.result()))

    print('Time taken for 1 epoch: {} secs\n'.format(time.time() - start))
    model_name = 'image_caption_transformer_' + str(epoch + 1) + '.h5'
    transformer.save_weights(model_name)

想试试Keras的ModelCheckpoint，但不知道怎么在自定义train_step的流程里用，有没有办法基于现有代码保存最佳模型，或者修改代码适配model.fit()？

方案一：在自定义训练循环中手动实现最佳模型保存 #

不用大改现有代码，只需要跟踪最佳性能指标（更建议加验证集用验证指标，避免过拟合，这里先以训练准确率为例），每个epoch结束后对比当前指标和历史最佳，只有当当前更好时才保存权重。

修改后的代码示例：

# 初始化最佳准确率，初始设为0
best_acc = 0.0
# 如果用损失作为判断标准，初始设为极大值：best_loss = float('inf')

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
    name='train_accuracy')

transformer = Transformer(num_layer, d_model, num_heads, dff, row_size, col_size, target_vocab_size,
                          max_pos_encoding=target_vocab_size, rate=dropout_rate)

@tf.function
def train_step(img_tensor, tar):
    tar_inp = tar[:, :-1]
    tar_real = tar[:, 1:]

    dec_mask = create_masks_decoder(tar_inp)

    with tf.GradientTape() as tape:
        predictions, _ = transformer(img_tensor, tar_inp,
                                     True,
                                     dec_mask)
        loss = loss_function(tar_real, predictions)

    gradients = tape.gradient(loss, transformer.trainable_variables)
    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))

    train_loss(loss)
    train_accuracy(tar_real, predictions)

for epoch in range(30):
    start = time.time()

    train_loss.reset_states()
    train_accuracy.reset_states()

    for (batch, (img_tensor, tar)) in enumerate(dataset):
        train_step(img_tensor, tar)

        if batch % 50 == 0:
            print('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(
                epoch + 1, batch, train_loss.result(), train_accuracy.result()))

    current_acc = train_accuracy.result().numpy()
    current_loss = train_loss.result().numpy()
    
    print('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1,
                                                        current_loss,
                                                        current_acc))

    # 对比准确率，更新并保存最佳模型
    if current_acc > best_acc:
        best_acc = current_acc
        print(f"新最佳准确率: {best_acc:.4f}, 保存模型...")
        transformer.save_weights('best_image_caption_transformer.h5')
    
    # 如果用损失判断，替换成这段：
    # if current_loss < best_loss:
    #     best_loss = current_loss
    #     print(f"新最佳损失: {best_loss:.4f}, 保存模型...")
    #     transformer.save_weights('best_image_caption_transformer.h5')

    print('单轮训练耗时: {} 秒\n'.format(time.time() - start))

方案二：修改代码适配model.fit()使用ModelCheckpoint #

把Transformer包装成Keras的自定义Model类，重写train_step方法，这样就能直接用model.fit()，并搭配ModelCheckpoint回调自动保存最佳模型。

代码示例：

class ImageCaptionTransformer(tf.keras.Model):
    def __init__(self, num_layer, d_model, num_heads, dff, row_size, col_size, target_vocab_size, dropout_rate):
        super().__init__()
        self.transformer = Transformer(num_layer, d_model, num_heads, dff, row_size, col_size, target_vocab_size,
                                      max_pos_encoding=target_vocab_size, rate=dropout_rate)
    
    def train_step(self, data):
        img_tensor, tar = data
        tar_inp = tar[:, :-1]
        tar_real = tar[:, 1:]

        dec_mask = create_masks_decoder(tar_inp)

        with tf.GradientTape() as tape:
            predictions, _ = self.transformer(img_tensor, tar_inp,
                                             training=True,
                                             look_ahead_mask=dec_mask)
            loss = self.compiled_loss(tar_real, predictions)

        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        self.compiled_metrics.update_state(tar_real, predictions)
        return {m.name: m.result() for m in self.metrics}

# 初始化模型
model = ImageCaptionTransformer(num_layer, d_model, num_heads, dff, row_size, col_size, target_vocab_size, dropout_rate)

# 编译模型，指定优化器、损失和指标
model.compile(
    optimizer=optimizer,
    loss=loss_function,
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]
)

# 设置ModelCheckpoint回调，保存最佳模型
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath='best_model.h5',
    monitor='sparse_categorical_accuracy',  # 对应你用的指标
    save_best_only=True,
    save_weights_only=True,  # 和你之前只保存权重的操作一致
    mode='max',  # 准确率越大越好，用损失的话改成'min'
    verbose=1
)

# 开始训练
model.fit(
    dataset,
    epochs=30,
    callbacks=[checkpoint_callback]
)

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