火车上分开并按组进行测试

Question

我有一个样本数据如下：

import pandas as pd

df = pd.DataFrame({"x": [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120],
                   "id": [1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5],
                   "label": ["a", "a", "a", "b", "a", "b", "b", "b", "a", "b", "a", "b"]})

所以我的数据看起来像这样

  x   id   label
 10   1    a
 20   1    a
 30   1    a
 40   1    b
 50   2    a
 60   2    b
 70   3    a
 80   3    a
 90   4    b
100   4    a
110   5    b
120   5    a

我想根据给定测试样本数量（例如 6 个样本）的标签分布，将此数据分为两组（训练、测试）。我的设置更喜欢将测试集的大小定义为表示测试样本数量的整数而不是百分比。但是，对于我的特定域，任何 id 都必须分配到仅一个组中。例如，如果将id为1的样本分配给训练集，则其他id为1的样本就不能分配给测试集。因此预期输出是 2 个数据帧，如下所示：

训练集

  x   id   label
 10   1    a
 20   1    a
 30   1    a
 40   1    b
 50   2    a
 60   2    b

测试集

  x   id   label
 70   3    a
 80   3    a
 90   4    b
100   4    a
110   5    b
120   5    a

训练集和测试集具有相同的类别分布（a:b 为 4:2），并且 id 1、2 仅分配给训练集，而 id 3、4、5 仅分配给测试集。我曾经用 sklearn 做过，train_test_split但我不知道如何在这种情况下应用它。我可以请教您如何处理此类情况的建议吗？

Answer 1

sklearn.model_selection除了 之外还有其他几个选项train_test_split。其中之一旨在解决您所追求的问题。在这种情况下，您可以使用GroupShuffleSplit，正如文档中提到的，它提供随机训练/测试索引来根据第三方提供的组分割数据。对于这些情况，您也有GroupKFold非常有用的。

from sklearn.model_selection import GroupShuffleSplit

X = df.drop('label',1)
y=df.label

您现在可以实例化GroupShuffleSplit，并像使用 一样执行操作train_test_split，唯一的区别是指定一个group列，该列将用于拆分X，y因此组将根据组值进行拆分：

gs = GroupShuffleSplit(n_splits=2, test_size=.6, random_state=0)
train_ix, test_ix = next(gs.split(X, y, groups=X.id))

现在您可以索引数据帧来创建训练集和测试集：

X_train = X.loc[train_ix]
y_train = y.loc[train_ix]

X_test = X.loc[test_ix]
y_test = y.loc[test_ix]

给予：

print(X_train)

      x  id
4    50   2
5    60   2
8    90   4
9   100   4
10  110   5
11  120   5

对于测试集：

print(X_test)

   x  id
0  10   1
1  20   1
2  30   1
3  40   1
6  70   3
7  80   3