在sklearn中使用标签拆分训练和测试集？答案

【问题标题】：Splitting train and test set with labels in sklearn?在sklearn中使用标签拆分训练和测试集？
【发布时间】：2018-12-26 16:46:05
【问题描述】：

我有一个扩充数据集，其中有一个列表 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]（具有相同标签）拆分为同一个子集。

【问题讨论】：

[3,3], [4,4] 一样吗？
@seralouk 是的，与 [5,5] 和 [6,6] 相同
查看我的回答，让我知道这是否适合您

标签： python python-3.x numpy scikit-learn

【解决方案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]

【讨论】：

为了将来参考，n_splits=3 是任意的，因为我们只希望获得一个折叠。