train_test_split#

sklearn.model_selection.train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)[来源]#

将数组或矩阵分割成随机的训练集和测试集子集。

一个快速的实用工具,它封装了输入验证、next(ShuffleSplit().split(X, y)),并将这些应用于输入数据,从而实现一行代码的数据分割(和可选的子采样)。

用户指南中阅读更多。

参数:
*arrays具有相同长度/形状[0]的可索引序列

允许的输入是列表、numpy数组、scipy稀疏矩阵或pandas数据帧。

test_size浮点数或整数,默认为 None

如果为浮点数,应介于 0.0 和 1.0 之间,表示要包含在测试分割中的数据集比例。如果为整数,则表示测试样本的绝对数量。如果为 None,则该值设置为训练集大小的补集。如果train_size也为 None,则它将被设置为 0.25。

train_size浮点数或整数,默认为 None

如果为浮点数,应介于 0.0 和 1.0 之间,表示要包含在训练分割中的数据集比例。如果为整数,则表示训练样本的绝对数量。如果为 None,则该值会自动设置为测试集大小的补集。

random_state整数、RandomState 实例或 None,默认为 None

控制在应用分割之前应用于数据的洗牌。传入一个整数以在多次函数调用中获得可重现的输出。参见词汇表

shuffle布尔值,默认为 True

是否在分割数据之前打乱数据。如果 shuffle=False 则 stratify 必须为 None。

stratify类数组,默认为 None

如果不为 None,则数据以分层方式分割,并将其用作类别标签。在用户指南中阅读更多。

返回:
splitting列表,长度 = 2 * len(arrays)

包含输入训练集-测试集分割的列表。

0.16 版新增:如果输入是稀疏的,输出将是 scipy.sparse.csr_matrix。否则,输出类型与输入类型相同。

示例

>>> import numpy as np
>>> from sklearn.model_selection import train_test_split
>>> X, y = np.arange(10).reshape((5, 2)), range(5)
>>> X
array([[0, 1],
       [2, 3],
       [4, 5],
       [6, 7],
       [8, 9]])
>>> list(y)
[0, 1, 2, 3, 4]
>>> X_train, X_test, y_train, y_test = train_test_split(
...     X, y, test_size=0.33, random_state=42)
...
>>> X_train
array([[4, 5],
       [0, 1],
       [6, 7]])
>>> y_train
[2, 0, 3]
>>> X_test
array([[2, 3],
       [8, 9]])
>>> y_test
[1, 4]
>>> train_test_split(y, shuffle=False)
[[0, 1, 2], [3, 4]]
>>> from sklearn import datasets
>>> iris = datasets.load_iris(as_frame=True)
>>> X, y = iris['data'], iris['target']
>>> X.head()
    sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)
0                5.1               3.5                1.4               0.2
1                4.9               3.0                1.4               0.2
2                4.7               3.2                1.3               0.2
3                4.6               3.1                1.5               0.2
4                5.0               3.6                1.4               0.2
>>> y.head()
0    0
1    0
2    0
3    0
4    0
...
>>> X_train, X_test, y_train, y_test = train_test_split(
... X, y, test_size=0.33, random_state=42)
...
>>> X_train.head()
    sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)
96                 5.7               2.9                4.2               1.3
105                7.6               3.0                6.6               2.1
66                 5.6               3.0                4.5               1.5
0                  5.1               3.5                1.4               0.2
122                7.7               2.8                6.7               2.0
>>> y_train.head()
96     1
105    2
66     1
0      0
122    2
...
>>> X_test.head()
    sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)
73                 6.1               2.8                4.7               1.2
18                 5.7               3.8                1.7               0.3
118                7.7               2.6                6.9               2.3
78                 6.0               2.9                4.5               1.5
76                 6.8               2.8                4.8               1.4
>>> y_test.head()
73     1
18     0
118    2
78     1
76     1
...