时间戳

首页 > 解决方案 > 在火车上拆分并按组分开测试

python - 在火车上拆分并按组分开测试

问题描述

我有一个样本数据如下：

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

标签： pythonpandasmachine-learningscikit-learn

解决方案

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

推荐阅读