XGBoost数据训练小例子

淩亂°似流年 2022-05-31 10:12 380阅读 0赞

# 1. 前言 #

XGBoost是很好的数据处理工具，可以在各大赛事中见到它的影子。本篇博客就主要针对对XGBoost的原理、相关Python API较为熟悉等的前提下将这些分散的内容串起来，从数据生成（已经准备好的数据，忽略数据预处理、特征选择之类的-\_-||）、训练模型、预测测试数据进行说明。希望能够对刚入手XGBoost的朋友有所帮助。

# 2. 完整过程 #

## 2.1 步骤分解 ##

1.  训练数据生成：这里直接采用的是生成标准的训练数据，其它的操作不考虑-\_-||。本例子中生成的是两个服从高斯分布的数据集
2.  训练模型：调用xgboost.train()实现
3.  数据预测：调用xgboost.predict()实现

## 2.2 完整代码 ##

# -*- coding=utf-8 -*-
    import numpy as np
    import matplotlib.pyplot as plt
    import xgboost as xgb
    
    
    # train data
    def get_train_data(data_size=100):
        data_label = np.zeros((2*data_size, 1))
        # class 1
        x1 = np.reshape(np.random.normal(3, 1, data_size), (data_size, 1))
        y1 = np.reshape(np.random.normal(4, 1, data_size), (data_size, 1))
        data_train = np.concatenate((x1, y1), axis=1)
        data_label[0:data_size, :] = 0
        # class 2
        x2 = np.reshape(np.random.normal(1, 1, data_size), (data_size, 1))
        y2 = np.reshape(np.random.normal(0.5, 1, data_size), (data_size, 1))
        data_train = np.concatenate((data_train, np.concatenate((x2, y2), axis=1)), axis=0)
        data_label[data_size:2*data_size, :] = 1
    
        return data_train, data_label
    
    
    # test data
    def get_test_data(start, end, data_size=100):
        data1 = (end - start) * np.random.random((data_size, 1)) + start
        data2 = (end - start) * np.random.random((data_size, 1)) + start
        data_test = np.concatenate((data1, data2), axis=1)
        return data_test
    
    
    # show data distribution
    def plot_data(train_data, data_size, test_data):
        plt.figure()
        plt.plot(train_data[0:data_size, 0], train_data[0:data_size, 1], 'g.',
                 train_data[data_size:2*data_size, 0], train_data[data_size:2*data_size, 1], 'b*',
                 test_data[:, 0], test_data[:, 1], 'rs')
        plt.legend(['class1', 'class 2', 'test_data'])
        plt.title('Distribution')
        plt.grid(True)
        plt.xlabel('axis1')
        plt.ylabel('axis2')
        plt.show()
    
    
    # plot predict res
    def plot_predict_data(train_data, data_size, test_data, predict_res1, predict_res2):
        plt.figure()
        plt.subplot(1, 2, 1)
        plt.plot(train_data[0:data_size, 0], train_data[0:data_size, 1], 'g.',
                 train_data[data_size:2*data_size, 0], train_data[data_size:2*data_size, 1], 'b*',
                 test_data[:, 0], test_data[:, 1], 'ms')
        plt.legend(['class1', 'class2', 'test_data'])
        plt.title('Distribution')
        plt.grid(True)
        plt.xlabel('axis1')
        plt.ylabel('axis2')
    
        plt.subplot(1, 2, 2)
        plt.plot(train_data[0:data_size, 0], train_data[0:data_size, 1], 'g.',
                 train_data[data_size:2 * data_size, 0], train_data[data_size:2 * data_size, 1], 'b*',
                 predict_res1[:, 0], predict_res1[:, 1], 'ro',
                 predict_res2[:, 0], predict_res2[:, 1], 'rs')
        plt.legend(['class1', 'class2', 'predict1', 'predict2'])
        plt.title('Predict res')
        plt.grid(True)
        plt.xlabel('axis1')
        plt.ylabel('axis2')
        plt.show()
    
    # main function
    if __name__ == '__main__':
        data_size = 100
        train_data0, label_data = get_train_data(data_size)  # train data generate
        test_data0 = get_test_data(-1, 5, 10)  # test data
        plot_data(train_data0, data_size, test_data0)  # plot
        # data convert
        train_data = xgb.DMatrix(train_data0, label=label_data)
        test_data = xgb.DMatrix(test_data0)
    
        # data training
        num_round = 50
        param = {
       'booster': 'gbtree', 'eta': 0.1, 'max_depth': 5, 'objective': 'binary:logistic'}
        bst = xgb.train(param, train_data, num_round)
    
        # make prediction
        predict_res = bst.predict(test_data)
        print predict_res
        index1 = predict_res > 0.5
        res1 = test_data0[index1, :]
        res2 = test_data0[~index1, :]
    
        # plot prediction result
        plot_predict_data(train_data0, data_size, test_data0, res1, res2)

## 2.3 结果 ##

在左图中是训练数据和测试数据的分布；右边是预测得到的结果，分别用红色实心圆和红色实心矩形表示不同的类别划分。  
![结果][70]

[70]: /images/20220531/141ef166942346409d5b566364919e5f.png