通过源代码注释了解sklearn中train_test_split的使用
sklearn中的train_test_split用于对数据集进行分割。如果不看文档,网上目前的教程主要都是将属性和标签分别进行分割,即:将 X 和 y 划分为 X_train, X_test, y_train, y_test 。事实上,该函数可以分割任意多的数据集,以更好地满足我们使用的需要。
首先,安装sklearn包并导入
from sklearn.model_selection import train_test_split
函数头格式为:
train_test_split(*arrays, **options)
可见,第一个参数为需要被分割的数据集,而第二个参数是一些选项。
源码中注释摘录与介绍
在函数的源代码中,详细地描述了各个参数的使用方法,并给出了例子,这里把源代码中的注释部分摘录,并进行了部分解释:
参数说明
Parameters
----------
*arrays : sequence of indexables with same length / shape[0]
Allowed inputs are lists, numpy arrays, scipy-sparse
matrices or pandas dataframes.
test_size : float or int, default=None
If float, should be between 0.0 and 1.0 and represent the proportion
of the dataset to include in the test split. If int, represents the
absolute number of test samples. If None, the value is set to the
complement of the train size. If ``train_size`` is also None, it will
be set to 0.25.
train_size : float or int, default=None
If float, should be between 0.0 and 1.0 and represent the
proportion of the dataset to include in the train split. If
int, represents the absolute number of train samples. If None,
the value is automatically set to the complement of the test size.
random_state : int or RandomState instance, default=None
Controls the shuffling applied to the data before applying the split.
Pass an int for reproducible output across multiple function calls.
See :term:`Glossary <random_state>`.
shuffle : bool, default=True
Whether or not to shuffle the data before splitting. If shuffle=False
then stratify must be None.
stratify : array-like, default=None
If not None, data is split in a stratified fashion, using this as
the class labels.
其中,必须的参数只有 *array 。通过说明可知,train_test_split 可以接收任意数量的待分割数据集,条件是它们的长度、每条数据的形状要相同。所以,在分割数据集时,如果拿到的是属性和标签在一起的一个数组,也可以直接进行分割,例如:
data_train, data_test =train_test_split(data,test_size=0.1, random_state=666)
可以将数据集data按照 9:1 分割为训练集和测试集。
参数 test_size、train_size指定了测试集、训练集的大小在整个数据集中的比例。这两个参数的类型可以是float或int。若为float型,可以取值在0~1之间,表示占比;若为int型,则表示数据集的具体大小(数据个数)。如果二者只指定了其一,另一个参数会自动设为互补值(float:1-x,int:[length of dataset]-x)。
另外一个比较常用的参数是 random_state 。当值为0时,每次分割数据使用的随机数都会不同。如果像上例那样给出一个特定的种子,如果其他参数相同,每次都会得到相同的分割结果。
返回值说明
Returns
-------
splitting : list, length=2 * len(arrays)
List containing train-test split of inputs.
.. versionadded:: 0.16
If the input is sparse, the output will be a
``scipy.sparse.csr_matrix``. Else, output type is the same as the
input type.
返回分割好的数据集。返回数据集的数量,是 *array 中包含的数据集数量的两倍。即,如果像上例中那样只有一个data,将返回两个数据集,其中训练集在前,测试集在后,例中返回 data_train, data_test;如果像多数教程中那样分割 X, y ,则会返回X_train, X_test, y_train, y_test 。
示例
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]]