分层训练/测试拆分，保证在严重不平衡的数据集上包含小类

Question

我正在处理大规模、不平衡的数据集，我需要在其中选择一个分层的训练集。但是，即使数据集非常不平衡，我仍然需要确保至少每个标签类都至少包含在训练集中一次。 sklearns train_test_split 或 StratifiedShuffleSplit 不会“保证”这种包含。

这是一个例子：

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(100).reshape((50, 2))
y = [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,3,3,3,3,3,4,4]

X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=4, random_state=42, stratify=y)

print(X_train, y_train)

结果是

[[80 81]
 [48 49]
 [18 19]
 [30 31]] [2, 2, 1, 1]

因此标签类 3 和 4 不包括在此训练拆分中。给定绝对的 train_size=4，这两个类还不够大，不能包含在内。对于严格分层的拆分，这是正确的。 但是，对于较小的类，我至少需要确保算法“已经看到标签类”。因此，我需要对分层原则进行某种软化，并按比例包含较小的班级。 我已经编写了相当多的代码来实现这一点，它首先删除较小的类，然后按比例拆分分别处理它们。但是，当移除时，由于类数量/总大小的变化，这也会影响 train_test_split。

是否有任何简单的函数/算法来实现这种行为？

Answer 1

你检查过sklearn.model_selection.StratifiedKFold吗？尝试将n_folds 设置为小于或等于人口最少的类中的成员数。如果你有，那么我只能推荐使用来自imbalanced-learn 的欠采样/过采样方法。