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

Question

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

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

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

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 ]