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

问题描述

我正在处理大规模、不平衡的数据集，我需要选择分层训练集。然而，即使数据集严重不平衡，我仍然需要确保至少每个标签类在训练集中至少包含一次。 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,2,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,1]

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

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

解决方法

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