cancel
Showing results for 
Search instead for 
Did you mean: 
Machine Learning
Dive into the world of machine learning on the Databricks platform. Explore discussions on algorithms, model training, deployment, and more. Connect with ML enthusiasts and experts.
cancel
Showing results for 
Search instead for 
Did you mean: 

Save model from AutoML to MLflow in LightGBM flavor

dkxxx-rc
Visitor

I want to get the LightGBM built-in variable importance values from a model that was generated by AutoML.  That's not logged in the metrics by default - can I change a setting so that it will be logged?

More fundamentally:  what I'd really like is to modify the LightGBM notebook generated by AutoML so that it logs the model to MLflow with flavor.loader_module equal to mlflow.lightgbm.  By default, it logs with that parameter equal to mlflow.sklearn.

As a result, if I want to load that model back in for new predictions, I have to use model = mlflow.pyfunc.load_model(), and that loads the model in a generic format that doesn't include useful LightGBM things like model.feature_importances.

I want to be able to load the model back in via model = mlflow.lightgbm.load_model(), but currently that generates an error because the model wasn't saved in the right format for that.

Now, I know model = mlflow.lightgbm.load_model() succeeds on a different model that I originally saved in LightGBM flavor via mlflow.lightgbm.log_model().  But the AutoML notebook doesn't use load_model(), so I have to look further for a way to force LightGBM flavor.

In that vein, I did find a command pyfunc.add_to_model(mlflow_model, loader_module="mlflow.sklearn") in the notebook.  Sadly, changing it to loader_module="mlflow.lightgbm" had no discernable effect on the problem.  The model saved to MLflow still had flavor.loader_module equal to mlflow.sklearn.

0 REPLIES 0

Connect with Databricks Users in Your Area

Join a Regional User Group to connect with local Databricks users. Events will be happening in your city, and you won’t want to miss the chance to attend and share knowledge.

If there isn’t a group near you, start one and help create a community that brings people together.

Request a New Group