pleas上的sklearn train_test_split按多列分层

Question

我是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

Answer 1

如果您想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']])

如果您担心由于11and3和1and 之类的值而13产生的碰撞，并且两者都创建了 的连接值113，那么您可以在中间添加一些任意字符串：

df['bc'] = df['b'].astype(str) + "_" + df['c'].astype(str)

Answer 2

你得到重复的原因是因为train_test_split()最终将strata定义为你传递给参数的任何值的唯一值集stratify.由于分层是从两列定义的,因此一行数据可能代表多个层,因此采样可以选择同一行两次,因为它认为它是从不同的类中采样的.

该train_test_split()函数调用 StratifiedShuffleSplit,它使用 np.unique()的y(这是通过你传递什么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作为一个单独的类的来样？如果是后者,那就是你已经得到的.前者更复杂,而这不是train_test_split设定的.

您可能会发现嵌套分层抽样的讨论很有用.

Answer 3

您正在使用哪个版本的scikit-learn？您可以sklearn.__version__用来检查。

在0.19.0之前的版本中，scikit-learn无法正确处理二维分层。它已在0.19.0中修复。

它在describled 问题＃9044。

更新您的scikit-learn应该可以解决问题。如果您无法更新scikit学习，请在此处查看此提交历史记录以获取修复。