在sklearn中用标签拆分列车和测试集

Question

我有一个扩充数据集，其中有一个列表label，用于对原始数据及其扩充进行分组。在sklearn中是否有类似于train_test_split()的方法，但保证相同的标签不会出现在两个子集中？我能想到的一个类似的方法是LabelKFold()。

例如，我有x = [[1,1], [2,2], [3,3], [4,4], [5,5], [6,6]]、y = [0, 0, 1, 1, 1, 1]、label = [0, 0, 1, 1, 2, 2]。

我希望[1,1]和[2,2](具有相同标签(被拆分为相同的子集。

Answer 1

您可以使用GroupKFold来执行此操作。以下是如何进行66/33的单次拆分。

from sklearn.model_selection import GroupKFold
gkf = GroupKFold(n_splits=3)
train, test = next(gkf.split(X, y, groups=label))
X_train = X[train]
y_train = y[train]
X_test = X[test]
y_test = y[test]