在火车上拆分并按组进行测试分离

Posted 2021-04-11

我有如下示例数据：

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，但我不知道如何在这种情况下应用它。请问您如何处理这种情况？

答案

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

GroupKFold

您现在可以实例化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]

另一答案

添加到Yatu出色的答案中，如果愿意，您只能使用熊猫来分割数据，尽管最好使用他的答案中提出的内容。

print(X_train)

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