是否有可能在两列的基础上对一组列车进行分层测试？

Question

考虑一个包含两个列( text和label )的dataframe。我可以很容易地创建一个分层的火车测试分裂使用sklearn.model_selection.train_测试_拆分。我唯一需要做的就是设置我想要用于分层的列(在本例中是label)。

现在，考虑包含三个列( text、subreddit和label )的dataframe。我想使用label列进行分层的列车测试分割，但我也希望确保在subreddit列方面没有偏见。例如，测试集可能有更多来自subreddit X的注释，而火车集则没有。

我如何在Python中做到这一点？

Answer 1

一种选择是将两个变量的数组提供给stratify参数，后者也接受多维数组。以下是scikit文档中的描述：

分层数组一样，default=None如果没有，数据以分层的方式分割，使用它作为类标签。

这里有一个例子：

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

# create dummy data with unbalanced feature value distribution
X = pd.DataFrame(np.concatenate((np.random.randint(0, 3, 500), np.random.randint(0, 10, 500)),axis=0).reshape((500, 2)), columns=["text", "subreddit"])
y = pd.DataFrame(np.random.randint(0,2, 500).reshape((500, 1)), columns=["label"])

# split stratified to target variable and subreddit col
X_train, X_test, y_train, y_test = train_test_split(
    X, pd.concat([X["subreddit"], y], axis=1), stratify=pd.concat([X["subreddit"], y], axis=1))

# remove subreddit cols from target variable arrays
y_train = y_train.drop(["subreddit"], axis=1)
y_test = y_test.drop(["subreddit"], axis=1)

如您所见，拆分也被分层为subreddit：

用于subreddits的

列车数据共享(

)

X_train.groupby("subreddit").count()/len(X_train)

给出

text
subreddit   
0   0.232000
1   0.232000
2   0.213333
3   0.034667
4   0.037333
5   0.045333
6   0.056000
7   0.056000
8   0.048000
9   0.045333

用于subreddits的

测试数据共享(

)

X_test.groupby("subreddit").count()/len(X_test)

给出

text
subreddit   
0   0.232
1   0.240
2   0.208
3   0.032
4   0.032
5   0.048
6   0.056
7   0.056
8   0.048
9   0.048

当然，只有当您有足够的数据同时分层到subreddit和目标变量时，这才能起作用。否则，scikit学习将抛出一个异常。