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