如何在 Databricks 的实验中检索 model.pkl

我想从我训练的模型中检索泡菜,我知道它在我在 Databricks 的实验中的运行文件中。

似乎 mlflow.pyfunc.load_model 只能执行 predict 方法。

有直接访问泡菜的选项吗?

我还尝试使用 pickle.load(path) 在运行中使用路径(路径示例:dbfs:/databricks/mlflow-tracking/20526156406/92f3ec23bf614c9d934dd0195/artifacts/model/model.pkl)。

stack overflow How can I retrive the model.pkl in the experiment in Databricks
原文答案

接受的答案

我最近找到了可以通过以下两种方法来完成的解决方案:

1.在保存模型时使用自定义的预测功能(查看 databricks 文档以获取更多详细信息)。

示例databricks给予

class AddN(mlflow.pyfunc.PythonModel):

    def __init__(self, n):
        self.n = n

    def predict(self, context, model_input):
        return model_input.apply(lambda column: column + self.n)
# Construct and save the model
model_path = "add_n_model"
add5_model = AddN(n=5)
mlflow.pyfunc.save_model(path=model_path, python_model=add5_model)

# Load the model in `python_function` format
loaded_model = mlflow.pyfunc.load_model(model_path)

2.在下载文物时,加载模型人工制品:

from mlflow.tracking import MlflowClient

client = MlflowClient()

tmp_path = client.download_artifacts(run_id="0c7946c81fb64952bc8ccb3c7c66bca3", path='model/model.pkl')

f = open(tmp_path,'rb')

model = pickle.load(f)

f.close()

client.list_artifacts(run_id="0c7946c81fb64952bc8ccb3c7c66bca3", path="")

client.list_artifacts(run_id="0c7946c81fb64952bc8ccb3c7c66bca3", path="model")

答案:

作者头像

使用 frmwk 的原生 load_model() 方法(例如 sklearn.load_model())或 download_artifacts()