XGBoost 1.3.1
更新时间:2023-01-18
XGBoost
XGBoost框架下,自定义作业支持发布保存模型为pickle
和joblib
格式,并且在发布至模型仓库时需要选择相应的模型文件。使用下面代码进行模型训练时,训练程序可以自行加载数据,训练数据选择空文件夹即可。
pickle格式示例代码
Python
1# -*- coding:utf-8 -*-
2""" xgboost train demo """
3import xgboost as xgb
4import numpy as np
5def save_model(model):
6 """ save model with pickle format """
7 import pickle
8 with open('output/clf.pickle', 'wb') as f:
9 pickle.dump(model, f)
10def save_model_joblib(model):
11 """ save model with joblib format """
12 import joblib
13 joblib.dump(model, 'output/clf.pkl')
14def main():
15 """ main """
16 rawData = np.array([[2, 4], [3, 4], [1, 2], [4, 5], [7, 8]])
17 label = np.array([6, 7, 3, 9, 15])
18 dtrain = xgb.DMatrix(rawData, label=label)
19 deval = xgb.DMatrix(np.array([[3, 5], [3, 6]]), label=np.array([8, 9]))
20 param = {'max_depth': 2, 'eta': 1, 'silent': 1, 'objective': 'reg:linear'}
21 evallist = [(deval, 'eval'), (dtrain, 'train')]
22 num_round = 10
23 bst = xgb.train(param, dtrain, num_round, evallist)
24 dtest = xgb.DMatrix(np.array([[2, 4], [7, 8]]))
25 ypred = bst.predict(dtest)
26 print(ypred)
27 save_model_joblib(bst)
28if __name__ == '__main__':
29 main()