XGBoost模型的保存和调用-CFANZ编程社区

使用Pickle保存和调用模型

pickle 是 python 编程中比较标准的一个保存和调用模型的库，我们可以使用 pickle 和 open 函数的连用，来将我们的模型保存到本地。

import pickle
import xgboost as xgb
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error as MSE, r2_score

# 加载数据
data = load_boston()
X = data.data
y = data.target

# 划分数据
Xtrain,Xtest,Ytrain,Ytest = train_test_split(X,y,test_size=0.3,random_state=0)

# 设定参数，对模型进行训练
dtrain = xgb.DMatrix(Xtrain,Ytrain) 
param = {'silent':True
         ,'obj':'reg:linear'
         ,"subsample":1
         ,"eta":0.05
         ,"gamma":20
         ,"lambda":3.5
         ,"alpha":0.2
         ,"max_depth":4
         ,"colsample_bytree":0.4
         ,"colsample_bylevel":0.6
         ,"colsample_bynode":1}
num_round = 180

bst = xgb.train(param, dtrain, num_round)

# 保存模型
pickle.dump(bst, open("model_1.dat","wb"))

注意，open中我们往往使用w或者r作为读取的模式，但其实w与r只能用于文本文件，当我们希望导入的不是文本文件，而是模型本身的时候，我们使用"wb"和"rb"作为读取的模式。其中wb表示以二进制写入，rb表示以二进制读入！模型会自动保存在当前工作目录下。

# 注意，如果我们保存的模型是xgboost库中建立的模型，则导入的数据类型也必须是xgboost库中的数据类型
dtest = xgb.DMatrix(Xtest,Ytest)

# 导入模型
loaded_model = pickle.load(open("model_1.dat", "rb"))
print("Loaded model from: model_1.dat")

# 模型预测
ypreds = loaded_model.predict(dtest)

MSE(Ytest,ypreds)

r2_score(Ytest,ypreds)

使用Joblib保存和调用模型

Joblib 是 SciPy 生态系统中的一部分，它为 Python 提供保存和调用管道和对象的功能，处理 NumPy 结构的数据尤其高效，对于很大的数据集和巨大的模型非常有用。Joblib 与 pickle API非常相似， 模型也会自动保存在当前工作目录下。

import joblib
bst = xgb.train(param, dtrain, num_round)

# 保存模型
joblib.dump(bst,"model_2.dat")

# 调用模型
loaded_model = joblib.load("model_2.dat")

ypreds = loaded_model.predict(dtest)
MSE(Ytest, ypreds)

r2_score(Ytest,ypreds)

年轻人，继续加油哇！！！