Databricks中MLflow Artifacts训练数据未显示的问题咨询

你的代码存在两个关键问题，导致训练数据未在MLflow UI中显示：

1. 缺少os模块导入 #

代码中使用os.path.join拼接文件路径，但未导入os模块，这会触发NameError，直接导致训练数据保存和上传的代码块无法执行。

2. numpy数组拼接的维度不匹配问题 #

X_train是二维numpy数组，y_train是一维numpy数组，直接用pd.concat拼接会出现维度不匹配的问题，无法生成正确的CSV文件。

修正后的完整代码 #

import mlflow
import mlflow.sklearn
import pandas as pd
import matplotlib.pyplot as plt
import tempfile
import os  # 新增os模块导入
from numpy import savetxt
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_diabetes
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

db = load_diabetes()
X = db.data
y = db.target
X_train, X_test, y_train, y_test = train_test_split(X, y)

with mlflow.start_run():
    # 设置模型参数
    n_estimators = 100
    max_depth = 6
    max_features = 3

    # 创建并训练模型
    rf = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth, max_features=max_features)
    rf.fit(X_train, y_train)

    # 生成测试集预测结果
    predictions = rf.predict(X_test) 

    # 记录模型参数
    mlflow.log_param("num_trees", n_estimators)
    mlflow.log_param("maxdepth", max_depth)
    mlflow.log_param("max_feat", max_features)
 
    # 计算并记录评估指标
    mse = mean_squared_error(y_test, predictions)
    mlflow.log_metric("mse", mse)

    # 记录模型
    mlflow.sklearn.log_model(rf, "random-forest-model")

    # 记录训练数据
    # 将numpy数组转为DataFrame/Series，确保维度匹配
    train_df = pd.concat([
        pd.DataFrame(X_train, columns=db.feature_names),
        pd.Series(y_train, name="target")
    ], axis=1)
    
    with tempfile.TemporaryDirectory() as tmp:
        path = os.path.join(tmp, 'train.csv')
        train_df.to_csv(path, index=False)  # 避免保存索引列，让数据更整洁
        mlflow.log_artifacts(tmp)
  
    mlflow.end_run()

额外优化说明 #

• 添加index=False参数，避免将DataFrame的索引列写入CSV文件，让训练数据结构更清晰。
• 为特征列和目标列添加命名，方便后续查看训练数据时快速理解各字段含义。

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