train_test_split()，随机划分训练集和测试集的函数

最新推荐文章于 2025-05-11 16:13:31 发布

原创最新推荐文章于 2025-05-11 16:13:31 发布 · 1.8w 阅读

33 ·

CC 4.0 BY-SA版权

文章标签：

#sklearn #train_test_split() #stratify #参数

python 专栏收录该内容

66 篇文章

订阅专栏

本文详细介绍了sklearn库中model_selection模块的train_test_split()函数，特别是stratify参数的使用。该参数确保了训练集和测试集在类别分布上的平衡，尤其适用于类别不平衡的数据集。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

sklearn.model_selection.train_test_split()

from sklearn.model_selection import train_test_split
x_train,x_test,y_train,y_test=
				train_test_split(train_data,train_lable,test_size=0.3,random_state=o,stratify=train_lable)

注意：stratify一般用于非平衡数据按train_lable的比例分层，多用于分类

train_data：样本特征集
train_target：样本的标签集
test_size：样本占比，测试集占数据集的比重，如果是整数的话就是样本的数量
random_state：是随机数的种子。在同一份数据集上，相同的种子产生相同的结果，不同的种子产生不同的划分结果
x_train,y_train:构成了训练集
x_test,y_test：构成了测试集

stratify参数详解

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_data（数据）就是按照X中的比例分配

将stratify=Y_data（标签）就是按照y中的比例分配