MLflow PyTorch Lightning自动日志为何未自动记录训练/验证指标？

可能的原因及对应解决方案 #

1. 自动日志配置冲突

你同时调用了全局mlflow.autolog()和PyTorch专用的mlflow.pytorch.autolog()，这会导致框架特定的自动日志逻辑被全局配置覆盖，丢失了针对PyTorch Lightning的指标捕获规则。

解决办法：移除全局自动日志调用，仅保留PyTorch专用的自动日志配置，并确保在Trainer初始化前执行：

mlflow.set_experiment(experiment_name = "path_to_databricks_experiment")
# 仅保留PyTorch专用自动日志，合并log_input_examples参数
mlflow.pytorch.autolog(log_every_n_epoch=1, log_models=True, log_input_examples=True)

trainer = pl.Trainer(max_epochs = 2, default_root_dir="path_to_databricks_experiment", log_every_n_steps=1)

with mlflow.start_run():
    trainer.fit(model, data_module)

2. PyTorch Lightning损失返回格式不规范

MLflow的自动日志依赖Lightning的默认损失跟踪机制，若你的training_step/validation_step未按标准格式返回损失，自动日志无法识别。

解决办法：确保步骤方法返回损失张量，或包含loss键的字典：

def training_step(self, batch, batch_idx):
    x, y = batch
    pred = self(x)
    loss = self.loss_fn(pred, y)
    # 必须返回loss张量或含loss键的字典
    return loss

3. 版本兼容性不足

早期版本的MLflow（<1.26.0）或PyTorch Lightning（<1.5.0）对自动日志的支持不完善，仅能记录模型参数，无法自动捕获训练/验证损失。

解决办法：检查当前环境版本，若过低则升级：

import mlflow
import pytorch_lightning as pl
print(f"MLflow version: {mlflow.__version__}")
print(f"PyTorch Lightning version: {pl.__version__}")

在Databricks中可通过集群库配置页面升级对应依赖包。

4. Trainer日志配置被禁用或覆盖

若Trainer手动设置了enable_logging=False，或指定了自定义logger参数，会导致MLflow自动注册的日志钩子失效。

解决办法：确保Trainer未禁用日志，且不手动指定与MLflow冲突的Logger：

# 保留默认enable_logging=True，不手动指定logger参数
trainer = pl.Trainer(max_epochs=2, default_root_dir="path_to_databricks_experiment", log_every_n_steps=1)

核心原因总结 #

问题本质是自动日志配置冲突或版本兼容性不足：

• 全局与框架专用自动日志的重复调用会覆盖针对Lightning的指标捕获逻辑；
• 旧版本依赖包未实现对Lightning训练/验证步骤损失的自动监听，必须显式调用self.log才能触发记录。

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