Keras训练中如何不使用autolog逐epoch向MLflow保存loss、accuracy指标

不用autolog的实现非常简单，Keras本身提供了完整的训练过程钩子，两种方式都能拿到所有轮次的指标，按需选择即可。

方式1：直接使用fit返回的历史对象（零额外代码，最简便） #

model.fit()执行完会自动返回一个History对象，它的history属性是标准字典，已经按epoch顺序存好了所有compile阶段定义的损失、指标，包含训练集和验证集的对应值，完全不需要额外写回调。
示例代码：

# 模型编译、数据集准备逻辑省略
history = model.fit(
    x_train,
    y_train,
    epochs=30,
    validation_data=(x_val, y_val),
    batch_size=64
)

# 直接取出所有逐epoch指标
epoch_metrics = history.history

返回的epoch_metrics结构示例：

{
    'loss': [0.542, 0.321, 0.213, ...], # 逐epoch训练损失
    'accuracy': [0.782, 0.876, 0.921, ...], # 逐epoch训练准确率
    'val_loss': [0.498, 0.356, 0.287, ...], # 逐epoch验证损失
    'val_accuracy': [0.801, 0.862, 0.894, ...] # 逐epoch验证准确率
}

如果你配置了自定义指标，比如compile时写了metrics=['accuracy', tf.keras.metrics.Precision()]，字典里会自动多出precision、val_precision对应的键值，直接取用即可。

方式2：自定义回调（适合训练过程中实时处理指标的场景） #

如果你需要在每个epoch结束时实时把指标上报到MLflow、做自定义逻辑判断、存额外参数（比如当前学习率、梯度统计值），自己写个继承tf.keras.callbacks.Callback的类即可，灵活度远高于autolog。
示例代码：

import tensorflow as tf
import mlflow

class CustomMetricsLogger(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        # 提取当前epoch的指标
        train_loss = logs.get('loss')
        train_acc = logs.get('accuracy')
        val_loss = logs.get('val_loss')
        val_acc = logs.get('val_accuracy')

        # 手动上报指标到MLflow，完全不需要调用autolog
        mlflow.log_metrics({
            'train_loss': train_loss,
            'train_acc': train_acc,
            'val_loss': val_loss,
            'val_acc': val_acc
        }, step=epoch)

        # 可扩展任意自定义逻辑，比如保存最优权重、异常告警等
        if val_acc > getattr(self, 'best_val_acc', 0):
            self.best_val_acc = val_acc
            self.model.save_weights('best_model.h5')
            print(f"Epoch {epoch+1} 刷新最优验证准确率: {val_acc:.4f}, 权重已保存")

训练时把回调实例传给fit的callbacks参数即可生效：

metrics_logger = CustomMetricsLogger()
history = model.fit(
    x_train,
    y_train,
    epochs=30,
    validation_data=(x_val, y_val),
    batch_size=64,
    callbacks=[metrics_logger]
)

注意：不推荐每个epoch都在测试集上跑评估，既会拖慢训练速度，也容易造成数据泄露。测试集指标建议等所有训练结束后，调用model.evaluate(x_test, y_test)一次性计算即可。

指标可视化 #

拿到指标字典之后，直接用绘图库就能生成直观的变化曲线：

import matplotlib.pyplot as plt
import pandas as pd

epochs = range(1, len(epoch_metrics['loss']) + 1)

# 绘制损失变化曲线
plt.figure(figsize=(12, 5))
plt.subplot(121)
plt.plot(epochs, epoch_metrics['loss'], label='训练集损失')
plt.plot(epochs, epoch_metrics['val_loss'], label='验证集损失')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('损失变化曲线')
plt.legend()

# 绘制准确率变化曲线
plt.subplot(122)
plt.plot(epochs, epoch_metrics['accuracy'], label='训练集准确率')
plt.plot(epochs, epoch_metrics['val_accuracy'], label='验证集准确率')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('准确率变化曲线')
plt.legend()

plt.tight_layout()
plt.show()

# 指标本地持久化存储为CSV，方便后续分析
pd.DataFrame(epoch_metrics, index=epochs).to_csv('epoch_training_metrics.csv', encoding='utf-8')

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