如何使用 cloudpickle 將 model 保存到 databricks DBFS 文件夾並加載它？

Question

我構建了一個 model，我的目標是將 model 保存為 pickle，稍后加載它以進行評分。 現在，我正在使用這段代碼：

    #save model as pickle
    import cloudpickle
    pickled = cloudpickle.dumps(final_model)
    
    #load model
    cloudpickle.loads(pickled)
    Output: <econml.dml.causal_forest.CausalForestDML at 0x7f388e70c373>

我擔心的是，通過這種方法，model 將僅保存在 Databricks 筆記本的會話變量“pickled”中。

不過，我希望將 model 存儲在 DBFS 存儲中，因此我可以隨時提取此 model（即使我的筆記本 session 過期）以使其更健壯。

我該怎么做？

Answer 1

您需要使用/dbfs/<DBFS path>使用本地文件 API 保存 model（請注意社區版不支持）。

但實際上，我建議依靠MLflow 將 model 和所有必要的超參數記錄到 MLflow Model 注冊表中 - 然后您可以使用 model 名稱輕松加載 model 而不是使用路徑。 使用 MLflow，您可以獲得更多好處，例如跟蹤多個 model 版本、使用階段等。