我有一个数据集，我希望将其分成训练和测试，以便在测试集中有每个数据源（在“源”列中指定）和每个类（在“类”列中指定）的数据。我读到有关将参数stratifiy与scikitlearn的{{1}}函数一起使用的信息，但是如何在两列上使用它呢？
在scikit-learn的两列上使用分层训练训练测试数据
•
问答
yndllhm 回答：在scikit-learn的两列上使用分层训练训练测试数据
从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(),1].mean()
(0.23472375,0.234725)
对均值均等进行统计t-test：
from scipy.stats import ttest_ind
ttest_ind(train_Y[:,0],0])
Ttest_indResult(statistic=0.0,pvalue=1.0)
最后对条件表示法做同样的事情，以证明您确实达到了想要的目标：
train_Y[train_Y[:,0].astype("bool"),test_Y[test_Y[:,1].mean()
(0.43959149751221877,0.43958874554180793)