在火车上拆分并按组分开测试答案

【问题标题】：Split on train and test separating by group在火车上拆分并按组分开测试
【发布时间】：2020-08-03 20:03:30
【问题描述】：

我有一个样本数据如下：

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，但我不知道如何在这种情况下应用它。我可以就如何处理这种情况提出您的建议吗？

【问题讨论】：

我想到的唯一方法是通过 id 拆分训练/测试（将唯一的 id 拆分为 2 组），但这可能不会产生所需的数据拆分百分比，这本来就可以从我所看到的情况来看，无法解决这个问题。

标签： python pandas machine-learning scikit-learn

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

【讨论】：

非常感谢。它完美地工作。我想知道...如果我想分成 3 组怎么办？当我设置输入参数 n_splits=3 时，ValueError: not enough values to unpack (expected 3, got 2)
是的，在您想要迭代拆分的情况下。查看示例here@RatchainantThammasudjarit 很高兴它有帮助:)
有没有办法确保这是分层拆分？

【解决方案2】：

除了 Yatu 的精彩回答之外，如果您愿意，您可以只使用 pandas 拆分您的数据，尽管最好使用他的回答中提出的建议。

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"],
    }
)


TRAIN_TEST_SPLIT_PERC = 0.75
uniques = df["id"].unique()
sep = int(len(uniques) * TRAIN_TEST_SPLIT_PERC)
df = df.sample(frac=1).reset_index(drop=True) #For shuffling your data
train_ids, test_ids = uniques[:sep], uniques[sep:]
train_df, test_df = df[df.id.isin(train_ids)], df[df.id.isin(test_ids)]


print("\nTRAIN DATAFRAME\n", train_df)
print("\nTEST DATAFRAME\n", test_df)

【讨论】：