将xgboost.Booster的实例转换为实现scikit-learn API的模型

问题描述

我正在尝试使用mlflow保存模型，然后稍后加载以进行预测。

我正在使用xgboost.XGBRegressor模型及其sklearn函数.predict()和.predict_proba()进行预测，但事实证明mlflow不支持实现sklearn API，因此，以后从mlflow加载模型时，mlflow返回xgboost.Booster的实例，并且它不实现.predict()或.predict_proba()函数。

是否可以将xgboost.Booster转换回实现sklearn API函数的xgboost.sklearn.XGBRegressor对象？

解决方法

您是否尝试过将模型包装在自定义类中，使用mlflow.pyfunc.PythonModel进行记录并加载？我举了一个简单的示例，并在重新加载模型后正确显示了<class 'xgboost.sklearn.XGBRegressor'>作为类型。

示例：

import xgboost as xgb
xg_reg = xgb.XGBRegressor(...)

class CustomModel(mlflow.pyfunc.PythonModel):
    def __init__(self,xgbRegressor):
        self.xgbRegressor = xgbRegressor

    def predict(self,context,input_data):
        print(type(self.xgbRegressor))
        
        return self.xgbRegressor.predict(input_data)

# Log model to local directory
with mlflow.start_run():
     custom_model = CustomModel(xg_reg)
     mlflow.pyfunc.log_model("custome_model",python_model=custom_model)


# Load model back
from mlflow.pyfunc import load_model
model = load_model("/mlruns/0/../artifacts/custome_model")
model.predict(X_test)

输出：

<class 'xgboost.sklearn.XGBRegressor'>
[ 9.107417 ]

mlflow save scikit-learn xgbclassifier xgboost