Pyspark如何保存IndexToString标签并通过MLflow读取用于预测值转换

你可以通过MLflow实现标签的存储和读取，常用两种方案可选：

方案1：单独记录indexer.labels到MLflow #

MLflow支持在训练run中存储自定义参数，你可以直接把标签列表作为参数关联到模型版本，预测时随模型一起读取即可。

训练阶段代码 #

indexer = StringIndexer(inputCol = target_variable_name, outputCol = 'label').fit(df)
df = indexer.transform(df)

# 此处省略随机森林模型训练逻辑，假设训练得到模型rf_model

with mlflow.start_run() as run:
    # 记录标签列表，标签数量多的话可以转JSON字符串存储
    mlflow.log_param("class_labels", indexer.labels)
    # 存储Spark随机森林模型
    mlflow.spark.log_model(rf_model, "RandomForest_model")

预测阶段代码 #

# 注意：你原代码用mlflow.sklearn.load_model加载Spark模型是错误的，需要用对应接口
rf_model = mlflow.spark.load_model("models:/RandomForest_model/latest")
# 从模型关联的训练run中读取标签
run = mlflow.get_run(rf_model.metadata.run_id)
class_labels = run.data.params["class_labels"]

# 转换预测结果
labelConverter = IndexToString(inputCol="prediction", outputCol="predictedLabel",labels=class_labels)
predictions = labelConverter.transform(rf_model.transform(new_data))

方案2：组装Pipeline统一存储（更推荐） #

直接把StringIndexer、随机森林模型、IndexToString都放进Spark ML Pipeline，训练完成后把整个Pipeline存到MLflow，后续加载可以直接输出原始分类值，不需要单独处理标签。

训练阶段代码 #

from pyspark.ml import Pipeline

# 定义全链路处理节点
indexer = StringIndexer(inputCol = target_variable_name, outputCol = 'label').fit(df)
rf = RandomForestClassifier(labelCol="label", featuresCol="features")
label_converter = IndexToString(inputCol="prediction", outputCol="predictedLabel", labels=indexer.labels)

# 组装Pipeline
pipeline = Pipeline(stages=[indexer, rf, label_converter])
pipeline_model = pipeline.fit(df)

# 直接存储整个Pipeline
with mlflow.start_run():
    mlflow.spark.log_model(pipeline_model, "rf_pipeline")

预测阶段代码 #

pipeline_model = mlflow.spark.load_model("models:/rf_pipeline/latest")
# 直接转换新数据，结果中predictedLabel列就是原始分类值，无需额外转换
predictions = pipeline_model.transform(new_data)

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