sklearn数据集分割函数[StratifiedShuffleSplit、split]

最新推荐文章于 2025-05-07 22:40:54 发布

墨氲

最新推荐文章于 2025-05-07 22:40:54 发布

阅读量6.4k

点赞数 14

CC 4.0 BY-SA版权

分类专栏： Python数据处理及可视化文章标签： python sklearn

本文链接：https://blog.youkuaiyun.com/dss_dssssd/article/details/82818601

Python数据处理及可视化专栏收录该内容

14 篇文章

订阅专栏

本文介绍sklearn中StratifiedShuffleSplit模块的使用方法，通过实例演示如何保持训练集和测试集中特征比例的一致性，确保模型评估的准确性。

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

sklearn.model_selection.StratifiedShuffleSplit(n_splits=10, test_size=’default’, train_size=None, random_state=None)
将数据集划分为几组数据集/测试集。
split(X, y)
按照y的值将数据集分为训练集或测试集，保证训练集和测试集中各类y值所占的比例与原数据集相同。

考虑一个场景，房屋数据集中认为地区收入在房价中占很重要的作用，想分开的训练集和测试集中各种收入所占的比重与原数据集中所占的比重相同。
代码实现：

    from sklearn.model_selection import StratifiedShuffleSplit
    split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
    for train_index, test_index in split.split(housing, housing["income_cat"]):
        strat_train_set = housing.loc[train_index]
        strat_test_set = housing.loc[test_index]

通过例子解释一下：
假设housing["income_cat"]=[1,1,1,1,1,2,2,22,2,2,2,2,2,2]也就是4个1，10个2，1所占的比例为1/3，2占的比例为2/3；
则经过split.split(housing, housing["income_cat"])后，strat_train_set [“income_cat”]和strat_test_set[“income_cat”]中1和2，所占的比例相同，分别为1/2，和2/3.也就是说，strat_train_set [“income_cat”]中有4个1，8个2；而strat_test_set[“income_cat”]中有1个1，2个2