PyTorch Lightning训练/验证/测试日志合并及MLFlow同图展示问题

核心问题是trainer.fit()和trainer.test()默认会分开记录指标，即使同一run，测试阶段的指标会被视为独立序列。可以通过以下两种方式合并：

方法一：手动复用训练Run并指定日志Step #

在训练完成后，手动捕获测试指标并将其日志到训练结束的全局Step，这样测试指标会被追加到训练/验证指标的同一序列中：

import mlflow
from pytorch_lightning import Trainer

# 初始化训练器和日志器（根据你的配置调整）
trainer = Trainer(logger=MLFlowLogger(...))

# 执行训练
trainer.fit(model, datamodule)

# 获取训练结束的全局Step
last_train_step = trainer.global_step

# 复用当前MLflow Run，日志测试指标
current_run_id = mlflow.active_run().info.run_id
with mlflow.start_run(run_id=current_run_id):
    # 执行测试并捕获结果
    test_metrics = trainer.test(model, datamodule, verbose=False)[0]
    
    # 按统一命名规范日志测试指标，指定Step为训练最后一步
    for metric_key, metric_value in test_metrics.items():
        mlflow.log_metric(f"test_{metric_key}", metric_value, step=last_train_step)

方法二：自定义Callback自动处理 #

写一个自定义回调，在测试阶段自动将指标日志到训练结束的Step，无需手动处理：

import mlflow
from pytorch_lightning.callbacks import Callback

class TestMetricSyncCallback(Callback):
    def on_test_epoch_end(self, trainer, pl_module):
        # 获取训练最后一步的Step
        final_train_step = trainer.global_step
        # 获取测试阶段的指标
        test_metrics = trainer.callback_metrics
        
        # 复用当前Run日志指标
        with mlflow.start_run(run_id=mlflow.active_run().info.run_id):
            for name, value in test_metrics.items():
                # 统一指标命名，和训练/验证的命名对应（比如val_loss → test_loss）
                mlflow.log_metric(f"test_{name}", value.item(), step=final_train_step)

初始化训练器时加入该回调：

trainer = Trainer(
    logger=MLFlowLogger(...),
    callbacks=[TestMetricSyncCallback()]
)

trainer.fit(model, datamodule)
trainer.test(model, datamodule)

关键注意事项 #

• 保持指标命名统一：训练用train_loss/train_auc，验证用val_loss/val_auc，测试用test_loss/test_auc，MLflow会自动将同后缀的指标合并到同一图表。
• 必须复用训练阶段的MLflow Run，避免创建新Run导致指标分离。

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