Scikit 学习是否有平均值，可以根据多个变量进行分层拆分？答案

【问题标题】：Is there a mean with Scikit learn, to do a stratified split based on several variables?Scikit 学习是否有平均值，可以根据多个变量进行分层拆分？
【发布时间】：2019-08-20 07:08:36
【问题描述】：

我正在研究一个数据框，我注意到 3 个变量对于预测标签非常重要。因此，我想将我的数据帧拆分为测试和训练集中，但不是随机拆分，而是基于这 3 个变量进行分层拆分（以保持训练集中与原始数据帧中的分布相同）。已创建函数StratifiedShuffleSplit 来处理标签，因此如果我没记错的话，我只能指定一个变量而不是三个变量。任何人都可以帮助我吗？谢谢

【问题讨论】：

欢迎来到 SO；请花点时间阅读How to Ask 和How to create a Minimal, Complete, and Verifiable example。

标签： python machine-learning scikit-learn data-science training-data

【解决方案1】：

这个交叉验证对象是 StratifiedKFold 和 ShuffleSplit 的合并，它返回分层的随机折叠。通过保留每个类的样本百分比来进行折叠。

注意：与 ShuffleSplit 策略一样，分层随机拆分并不能保证所有折叠都不同，尽管这对于相当大的数据集仍然很有可能。

>>> StratifiedShuffleSplit(n_splits=5, random_state=0, ...)
>>> for train_index, test_index in sss.split(X, y):
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]

【讨论】：