使用sklearn.model_selection.train_test_split可以在数据集上随机划分出一定比例的训练集和测试集,并返回拆分得到的train和test数据集。
使用格式:
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(train_data, train_target, test_size = 0.25, random_state = 0,stratify = y_train)
参数解释:
train_data:被划分的样本特征集。
train_target:被划分的样本标签集。
test_size:如果是浮点数,且在0-1之间,表示样本占比,测试集占数据集的比重;如果是整数的话则为训练集样本的绝对数量。
random_state:随机数的种子。
随机数种子:其实就是该组随机数的编号,在需要重复试验的时候,保证得到一组一样的随机数。比如你每次都填1,其他参数一样的情况下你得到的随机数组是一样的。但填0或不填,每次都会不一样。
随机数的产生取决于种子,随机数和种子之间的关系遵从以下两个规则:
种子不同,产生不同的随机数;种子相同,即使实例不同也产生相同的随机数。
shuffle:布尔值。是否在拆分前重组数据。如果shuffle=False,则stratify必须为None。
stratify:array-like or None。如果不是None,则数据集以分层方式拆分,并使用此作为类标签。
stratify是为了保持split前类的分布。比如有100个数据,80个属于A类,20个属于B类。如果train_test_split(… test_size=0.25, stratify = y_all), 那么split之后数据如下:
training: 75个数据,其中60个属于A类,15个属于B类。
testing: 25个数据,其中20个属于A类,5个属于B类。
用了stratify参数,training集和testing集的类的比例是 A:B= 4:1,等同于split前的比例(80:20)。通常在这种类分布不平衡的情况下会用到stratify。将stratify=X就是按照X中的比例分配 将stratify=y就是按照y中的比例分配 。
示例代码:
>>> 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]]
参考链接:https://blog.youkuaiyun.com/jiushinayang/article/details/81098186
本文介绍如何使用sklearn的train_test_split函数在数据集上随机划分训练集和测试集,包括函数的使用格式、参数解释及示例代码。通过设置test_size、random_state和stratify等参数,可以灵活控制数据集的划分比例和随机性。
4628





