在 scikit-learn 中分层 Train/Test-split

Question

我需要将我的数据分成训练集 (75%) 和测试集 (25%)。我目前使用以下代码执行此操作：

X, Xt, userInfo, userInfo_train = sklearn.cross_validation.train_test_split(X, userInfo)

但是，我想对我的训练数据集进行分层。我怎么做？我一直在研究 StratifiedKFold 方法，但没有让我指定 75%/25% 的分割并且只对训练数据集进行分层。

Answer 1

[0.17 更新]

请参阅 sklearn.model_selection.train_test_split 的文档：

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    stratify=y, 
                                                    test_size=0.25)

[/0.17 更新]

有一个拉取请求 here。但你可以简单地做 train, test = next(iter(StratifiedKFold(...))) 并根据需要使用训练和测试索引。

Answer 2

TL;DR : 使用 StratifiedShuffleSplit 和 test_size=0.25

Scikit-learn 为分层拆分提供了两个模块：

StratifiedKFold ：此模块可用作直接 k 折交叉验证运算符：因为它将设置 n_folds training/testing 集，使得类两者都相当平衡。

这是一些代码（直接来自上面的文档）

>>> skf = cross_validation.StratifiedKFold(y, n_folds=2) #2-fold cross validation
>>> len(skf)
2
>>> for train_index, test_index in skf:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
...    #fit and predict with X_train/test. Use accuracy metrics to check validation performance

StratifiedShuffleSplit ：此模块创建一个 training/testing 集，具有相同的平衡（分层）类。本质上，这就是您想要的 n_iter=1。您可以在此处提及测试大小，与 train_test_split

代码：

>>> sss = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=0)
>>> len(sss)
1
>>> for train_index, test_index in sss:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
>>> # fit and predict with your classifier using the above X/y train/test

Answer 3

这是 continuous/regression 数据的示例（直到 this issue on GitHub 得到解决）。

min = np.amin(y)
max = np.amax(y)

# 5 bins may be too few for larger datasets.
bins     = np.linspace(start=min, stop=max, num=5)
y_binned = np.digitize(y, bins, right=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, 
    y, 
    stratify=y_binned
)

其中 start 是最小值，stop 是连续目标的最大值。
如果你不设置 right=True 那么它或多或少会让你的最大值成为一个单独的 bin 并且你的拆分总是会失败，因为在那个额外的 bin 中的样本太少。

Answer 4

除了@Andreas Mueller 已接受的答案之外，只想添加上面提到的@tangy：

StratifiedShuffleSplit most closely resembles train_test_split（分层=y）添加了以下功能：

默认分层
通过指定n_splits，重复拆分数据

Answer 5

#train_size is 1 - tst_size - vld_size
tst_size=0.15
vld_size=0.15

X_train_test, X_valid, y_train_test, y_valid = train_test_split(df.drop(y, axis=1), df.y, test_size = vld_size, random_state=13903) 

X_train_test_V=pd.DataFrame(X_train_test)
X_valid=pd.DataFrame(X_valid)

X_train, X_test, y_train, y_test = train_test_split(X_train_test, y_train_test, test_size=tst_size, random_state=13903)

Answer 6

您可以简单地使用 Scikit learn 中提供的 train_test_split() 方法来完成：

from sklearn.model_selection import train_test_split 
train, test = train_test_split(X, test_size=0.25, stratify=X['YOUR_COLUMN_LABEL'])

我还准备了一个简短的 GitHub 要点，展示了 stratify 选项的工作原理：

https://gist.github.com/SHi-ON/63839f3a3647051a180cb03af0f7d0d9

Answer 7

正在将上面的@tangy 回答更新到 scikit-learn 的当前版本：0.23.2 (StratifiedShuffleSplit documentation).

from sklearn.model_selection import StratifiedShuffleSplit

n_splits = 1  # We only want a single split in this case
sss = StratifiedShuffleSplit(n_splits=n_splits, test_size=0.25, random_state=0)

for train_index, test_index in sss.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

Answer 8

因此，希望将数据集拆分为训练集和测试集，以便在每个 class 中保留与原始数据集中观察到的示例相同的比例。

这称为分层训练-测试拆分。

我们可以通过将“stratify”参数设置为原始数据集的 y 分量来实现这一点。 train_test_split() 函数将使用它来确保训练集和测试集在提供的“y”数组中存在的每个 class 中的示例比例。

Answer 9

StratifiedShuffleSplit 在我们选择应该在我们将要生成的所有小数据集中均匀表示的列之后完成。 'The folds are made by preserving the percentage of samples for each class.'

假设我们有一个包含 'season' 列的数据集 'data'，我们希望得到 'season' 的偶数表示，那么它看起来像这样：

from sklearn.model_selection import StratifiedShuffleSplit
sss=StratifiedShuffleSplit(n_splits=1,test_size=0.25,random_state=0)

for train_index, test_index in sss.split(data, data["season"]):
    sss_train = data.iloc[train_index]
    sss_test = data.iloc[test_index]

在 scikit-learn 中分层 Train/Test-split

Stratified Train/Test-split in scikit-learn

python

scikit-learn