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

Answer 2

除了Yatu出色的答案外，您也可以只用熊猫来拆分数据，尽管最好使用他的答案中提出的建议。

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)

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

2 个答案: