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()