python-3.x - 如何根据组 ID 生成训练-测试-分割？

Question

我有以下数据:

pd.DataFrame({'Group_ID':[1,1,1,2,2,2,3,4,5,5],
          'Item_id':[1,2,3,4,5,6,7,8,9,10],
          'Target': [0,0,1,0,1,1,0,0,0,1]})

   Group_ID Item_id Target
0         1       1      0
1         1       2      0
2         1       3      1
3         2       4      0
4         2       5      1
5         2       6      1
6         3       7      0
7         4       8      0
8         5       9      0
9         5      10      1

我需要根据“Group_ID”将数据集拆分为训练集和测试集，以便 80% 的数据进入训练集，20% 进入测试集。

也就是说，我需要我的训练集看起来像:

    Group_ID Item_id Target
0          1       1      0
1          1       2      0
2          1       3      1
3          2       4      0
4          2       5      1
5          2       6      1
6          3       7      0
7          4       8      0

和测试集:

Test Set
   Group_ID Item_id Target
8         5       9      0
9         5      10      1

最简单的方法是什么？据我所知，sklearn 中的标准 test_train_split 函数不支持按组拆分，而我也可以指示拆分的大小(例如 80/20)。

Answer 1

我找到了答案。这似乎有效:

from sklearn.model_selection import GroupShuffleSplit 

splitter = GroupShuffleSplit(test_size=.20, n_splits=2, random_state = 7)
split = splitter.split(df, groups=df['Group_Id'])
train_inds, test_inds = next(split)

train = df.iloc[train_inds]
test = df.iloc[test_inds]