mlflow.pyfunc.spark_udf 和向量结构类型

Question

我的 PySpark 数据集包含分类数据。

为了根据这些数据训练模型，我关注了example notebook。特别是，请参阅编码部分的预处理数据部分。

我现在需要在其他地方使用这个模型；因此，我遵循 Databricks 的建议来保存和加载此模型。

Pandas 可以正常工作（参见下面的代码）。

logged_model = 'runs:/e905f5759d434a1391bbe1e54a2b/best-model'

# Load model as a PyFuncModel.
loaded_model = mlflow.pyfunc.load_model(logged_model)

# Predict on a Pandas DataFrame.
import pandas as pd
loaded_model.predict(pd.DataFrame(data))

但是数据框太大，无法转换为 Pandas。因此我需要让它在 Spark 中工作：

import mlflow
logged_model = 'runs:/e905f5759d434a131bbe1e54a2b/best-model'

# Load model as a Spark UDF.
loaded_model = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model)

# Predict on a Spark DataFrame.
df.withColumn('predictions', loaded_model(*columns)).collect()

但是这个 sn-p 正在生产：

java.lang.UnsupportedOperationException: Unsupported data type: struct<type:tinyint,size:int,indices:array<int>,values:array<double>>

我的感觉是 udf 不接受这种类型的数据作为输入。 有没有办法解决它？ 另一种解决方案？

Answer 1

您是否尝试过使用mlflow.spark.load_model？

我在这里遇到了一个非常相似的问题，但使用的是 spark 方法。我尝试使用mlflow.spark.load_model('runs:/run-id/my-model') 方法，我得到了这个奇怪的错误：

FileNotFoundError: [Errno 2] No such file or directory: '/dbfs/tmp/mlflow/weird-id-folder'

搜索文档，我看到我们面临的问题（似乎不同），似乎是一个签名问题。

根据docs 的其他部分，我们知道模型记录的签名将有助于定义模型的输入类型。对我来说，问题是我的输入是 Spark 稀疏向量——不支持它......现在我正在尝试将其转换为基于列的签名。

你尝试过这样的事情吗？

---

更新：

我想补充一点，在我的情况下，添加签名确实解决了问题。我所做的只是忽略向量，只考虑输入数据和输出数据。

我查看了笔记本，但没有看到任何 mlflow 日志，无论如何，我认为您正在根据this 并使用mlflow.spark 风格记录您的实验。

如果是这样，请考虑使用 from pyspark.ml import Pipeline 在同一管道中使用所有数据转换和模型拟合。在记录模型之前，请考虑进行签名并注册模型架构。

import mlflow.spark
from mlflow.models.signature import infer_signature

with mlflow.start_run():
    [...]
    # executing train & test pipelines:
    model = pipeline.fit(train_features) # training model
    predictions = model.transform(test_features) # testing model
    train_signature = train_features.select('input_data') # ignores all other features created on the pipeline
    prediction_signature = predictions.select('input_data', 'prediction') # ignores all other features created on the training pipeline 
    signature = infer_signature(train_signature, prediction_signature) # register model schema
    mlflow.spark.log_model(model, 'transactions-classification', signature=signature) # logging model to mlflow
    [...]

将模型记录到实验后，在不同的笔记本中，您可以将 load_model 函数用作：

# importing model
import mlflow.spark
model_path = 'runs:/run-id'
model = mlflow.spark.load_model(model_path)

它会起作用的！ :D