在 scikit-learn 的两列上使用分层训练测试数据拆分答案

【问题标题】：train test data split using stratify on two columns in scikit-learn在 scikit-learn 的两列上使用分层训练测试数据拆分
【发布时间】：2020-06-22 02:10:25
【问题描述】：

我有一个数据集，我想将其拆分为训练和测试，以便在测试集中拥有来自每个数据源（在“源”列中指定）和每个类（在“类”列中指定）的数据。我阅读了有关将参数stratifiy 与scikitlearn 的train_test_split 函数一起使用的信息，但是如何在两列上使用它？

【问题讨论】：

您需要为此编写自己的包装器，目前此功能在 sklearn 中不可用。

标签： scikit-learn train-test-split

【解决方案1】：

自 v.19.0 起，使用 sklearn's train_test_split 可以轻松地在多个列上进行分层

证明

from sklearn.model_selection import train_test_split
from sklearn.datasets import make_multilabel_classification

X, Y = make_multilabel_classification(1000000, 10, n_classes=2, n_labels=1)
train_X, test_X, train_Y, test_Y =train_test_split(X,Y,stratify=Y, train_size=.8, random_state=42)
Y.shape

(1000000, 2)

然后您可以比较结果分层的简单列均值：

train_Y[:,0].mean(), test_Y[:,0].mean()
(0.45422, 0.45422)

train_Y[:,1].mean(), test_Y[:,1].mean()
(0.23472375, 0.234725)

对均值进行统计t-test：

from scipy.stats import ttest_ind
ttest_ind(train_Y[:,0],test_Y[:,0])

Ttest_indResult(statistic=0.0, pvalue=1.0)

最后对条件手段做同样的事情来证明你确实达到了你想要的：

train_Y[train_Y[:,0].astype("bool"),1].mean(), test_Y[test_Y[:,0].astype("bool"),1].mean()
(0.43959149751221877, 0.43958874554180793)

【讨论】：