利用sklearn划分训练集和测试集

利用sklearn划分训练集和测试集

  交叉验证（Cross Validation）是常用的机器学习训练手段，可以有效检验一个模型的泛化能力。交叉验证需要将原始数据集平等地划分为若干份，例如 5-folds CV 指的是将数据集分为5份，然后进行5次训练，每次取出一份数据作为测试集，剩下的作为训练集，得到5个模型，最终将5个模型的预测值做一个平均。
  CV的第一步就是划分数据集，sklearn.model_selection中的train_test_split提供了这个功能，传入数据即可返回划分好的数据。

train_test_split API

函数原型

sklearn.model_selection.train_test_split(*arrays, **options)

主要参数

• X ：数据集
• Y ：标签
• test_size ：用于作为测试集的比例，$\in (0,1)$
• random_state : 随机数种子

使用实例

  train_test_split只能传入一个比例将数据分为两部分，如果想要将数据分为若干部分，例如要将数据集分为四等分，可以先二等分数据集，然后递归地二等分两个子集，即可实现四等分。

# coding:utf-8
import pandas as pd
from sklearn.model_selection import train_test_split

data = pd.read_csv("./test.csv")
label = data['label']

# 首先二等分
data_half_1, data_half_2, label_half_1, label_half_2 = \
    train_test_split(data, label, test_size=0.5, random_state=19931028)

# 第一个子集二等分
data_qua_1, data_qua_2, label_qua_1, label_qua_2 = \
    train_test_split(data_half_1, label_half_1, test_size=0.5, random_state=19931028)

# 第二个子集二等分
data_qua_3, data_qua_4, label_qua_3, label_qua_4 = \
    train_test_split(data_half_2, label_half_2, test_size=0.5, random_state=19931028)

data_qua_1.to_csv('test_part1.csv', index=False) # part 1
data_qua_2.to_csv('test_part2.csv', index=False) # part 2
data_qua_3.to_csv('test_part3.csv', index=False) # part 3
data_qua_4.to_csv('test_part4.csv', index=False) # part 4 

注意事项

  经过观察，使用train_test_split进行分割后，并不保证每个子集中正负样本的比例与原数据集一致。

参考资料

【API】sklearn.model_selection.train_test_split