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

Question

我正在尝试使用 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 对象？

Answer 1

您是否尝试过在自定义 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 ]

Answer 2

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

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

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

scikit-learn

save

xgboost

mlflow

xgbclassifier