我是sklearn的一个相对较新的用户，并且在sklearn.model_selection中遇到了train_test_split中的一些意外行为。我有一个熊猫数据框，我想分成一个训练和测试集。我想将数据分层至少2个，但理想情况下我的数据框中有4列。
当我尝试这样做时，sklearn没有警告，但后来我发现我的最终数据集中有重复的行。我创建了一个示例测试来显示这种行为：
from sklearn.model_selection import train_test_split
a = np.array([i for i in range(1000000)])
b = [i%10 for i in a]
c = [i%5 for i in a]
df = pd.DataFrame({'a':a, 'b':b, 'c':c})
如果按任意一列分层，它似乎按预期工作：
train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['b']])
print(len(train.a.values)) # prints 800000
print(len(set(train.a.values))) # prints 800000
train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['c']])
print(len(train.a.values)) # prints 800000
print(len(set(train.a.values))) # prints 800000
但是当我尝试按两列进行分层时，我得到重复的值：
train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['b', 'c']])
print(len(train.a.values)) # prints 800000
print(len(set(train.a.values))) # prints 640000
答案 0 :(得分：10)
你得到重复的原因是因为train_test_split()最终将strata定义为传递给stratify参数的唯一值集。由于分层是从两列定义的，因此一行数据可能代表多个层，因此采样可以选择同一行两次，因为它认为它是从不同的类中采样。
train_test_split()函数calls StratifiedShuffleSplit，np.unique()上的uses y（您通过{{1}传递的内容}）。来自源代码：
stratify这是一个简化的示例案例，您提供的示例的变体：
classes, y_indices = np.unique(y, return_inverse=True)
n_classes = classes.shape[0]
分层函数认为有四个类可以拆分：from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
N = 20
a = np.arange(N)
b = np.random.choice(["foo","bar"], size=N)
c = np.random.choice(["y","z"], size=N)
df = pd.DataFrame({'a':a, 'b':b, 'c':c})
print(df)
a b c
0 0 bar y
1 1 foo y
2 2 bar z
3 3 bar y
4 4 foo z
5 5 bar y
...
，foo，bar和y。但由于这些类基本上是嵌套的，意味着z和y都显示在z和b == foo中，当分割器尝试从每个类中进行采样时，我们会得到重复的
b == bar这里有一个更大的设计问题：您是否想要使用嵌套分层抽样，或者您实际上只是想将train, test = train_test_split(df, test_size=0.2, random_state=0,
stratify=df[['b', 'c']])
print(len(train.a.values)) # 16
print(len(set(train.a.values))) # 12
print(train)
a b c
3 3 bar y # selecting a = 3 for b = bar*
5 5 bar y
13 13 foo y
4 4 foo z
14 14 bar z
10 10 foo z
3 3 bar y # selecting a = 3 for c = y
6 6 bar y
16 16 foo y
18 18 bar z
6 6 bar y
8 8 foo y
18 18 bar z
7 7 bar z
4 4 foo z
19 19 bar y
#* We can't be sure which row is selecting for `bar` or `y`,
# I'm just illustrating the idea here.
和df.b中的每个类视为一个单独的类来进行抽样？如果是后者，那就是你已经得到的。前者更复杂，而不是df.c设置要做的事情。
您可能会发现this discussion嵌套分层抽样有用。
答案 1 :(得分：4)
您使用的是什么版本的scikit-learn？您可以使用sklearn.__version__进行检查。
在版本0.19.0之前，scikit-learn无法正确处理二维分层。它修补了0.19.0。
在issue #9044中对其进行了描述。
更新你的scikit-learn应该解决问题。如果您无法更新scikit-learn，请参阅此修订提交历史记录here。
答案 2 :(得分：3)
如果您希望train_test_split的行为符合预期（通过没有重复的多个列进行分层），请创建一个新列，该列由其他列中的值组成，并在新列上进行分层。
df['bc'] = df['b'].astype(str) + df['c'].astype(str)
train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['bc']])
如果您担心由于11和3和1和13之类的值而产生的碰撞值113的冲突，那么您可以在中间添加任意字符串：
df['bc'] = df['b'].astype(str) + "_" + df['c'].astype(str)