train_test_split用法详解

格式

def train_test_split(*arrays,
                     test_size=None,
                     train_size=None,
                     random_state=None,
                     shuffle=True,
                     stratify=None)
参数详解:

在这里插入图片描述
关于stratify参数的进一步解释:

其作用是保持划分后的训练集与测试集中的数据分类的比例与原数据一致

例如,有100个数据,10个正例,90个反例,类别比例为1:10,经过train_test_split()划分后的训练集与验证集中数据的正例反例的比例不一定能保持1:10,而令stratify=标签数组时,便能够保持正反例的比例。

一般用法:

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1,stratify=y)
输出详解

在这里插入图片描述

实例
   '''
    Examples
    --------
    >>> 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]]
    '''

更多推荐