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

Convert an instance of xgboost.Booster into a model that implements the scikit-learn API

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

我正在使用 xgboost.XGBRegressor 模型及其 sklearn 函数 .predict().predict_proba() 进行预测,但事实证明 mlflow 不支持以下模型实现了 sklearn API,因此当稍后从 mlflow 加载模型时,mlflow returns 是 xgboost.Booster 的一个实例,并且它没有实现 .predict().predict_proba() 函数。

有没有办法将 xgboost.Booster 转换回实现 sklearn API 函数的 xgboost.sklearn.XGBRegressor 对象?

您是否尝试过在自定义 class 中包装您的模型,使用 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 ]

我有一个 xgboost.core.Booster 对象,它可以进行 return 概率计算,如下所示 your_Booster_model_object.predict(your_xgboost_dmatrix_dataset).