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

Question

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

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

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

Answer 1

您是否尝试过将模型包装在自定义类中，使用 记录和加载它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 ]