sklearn.model_selection.train_test_split函数

函数原型

sklearn.model_selection.train_test_split(
    *arrays,
    test_size=None,
    train_size=None,
    random_state=None,
    shuffle=True,
    stratify=None,
)

函数说明

train_test_split函数用于将数据集细分为训练集和测试集。参数*arrays表示待划分的数据集，可以是列表、numpy数组、scipy稀疏矩阵或pandas的数据。

参数test_size如果为浮点数，则表示测试集占数据集的百分比；如果为整数，则表示测试集的数量。参数train_size如果为浮点数，则表示训练集占数据集的百分比；如果为整数，则表示训练集的数量。

参数random_state表示随机数的种子，参数shuffle表示是否打乱数据。

参数stratify是为了保持split前类的分布。比如有100个数据，80个属于A类，20个属于B类。如果train_test_split(… test_size=0.25, stratify = y), 那么split之后数据如下：
training: 75个数据，其中60个属于A类，15个属于B类。
testing: 25个数据，其中20个属于A类，5个属于B类。

函数使用

>>> from sklearn import model_selection, datasets
>>> data, labels = datasets.load_iris(return_X_y=True)
>>> train_data, test_data = model_selection.train_test_split(data, test_size=0.4)
>>> len(train_data)
90
>>> len(test_data)
60
>>> train_labels, test_labels = model_selection.train_test_split(labels, test_size=0.4)
>>> train_labels
array([0, 1, 0, 0, 0, 0, 1, 2, 0, 0, 2, 2, 2, 2, 1, 0, 0, 1, 1, 2, 0, 0,
       2, 0, 1, 1, 1, 1, 0, 1, 2, 2, 1, 1, 0, 1, 2, 2, 2, 0, 2, 2, 1, 1,
       1, 1, 2, 1, 2, 0, 0, 2, 0, 0, 1, 0, 2, 2, 0, 1, 2, 0, 1, 0, 2, 2,
       0, 1, 0, 1, 1, 2, 1, 1, 1, 0, 2, 0, 1, 2, 1, 0, 0, 0, 1, 2, 0, 0,
       0, 0])
>>> train_labels, test_labels = model_selection.train_test_split(labels, test_size=0.4, random_state=10)
>>> train_labels
array([0, 0, 2, 1, 2, 0, 2, 0, 1, 1, 0, 2, 2, 2, 2, 2, 0, 1, 2, 1, 0, 2,
       1, 1, 0, 0, 0, 1, 2, 2, 1, 0, 0, 0, 2, 2, 1, 1, 2, 2, 2, 2, 1, 0,
       0, 1, 0, 0, 2, 1, 0, 0, 0, 1, 0, 1, 0, 1, 2, 0, 1, 1, 2, 0, 2, 0,
       1, 1, 2, 2, 0, 1, 2, 2, 1, 1, 2, 0, 2, 0, 0, 1, 0, 2, 2, 2, 1, 0,
       2, 0])
>>> train_labels, test_labels = model_selection.train_test_split(labels, test_size=0.4, random_state=10, shuffle=False)
>>> train_labels
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
       1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
       1, 1])