python - 使用 Numpy 分层将数据拆分为训练、测试、验证

Question

我刚刚看到了这个answer on SO它展示了如何使用 numpy 分割数据。

假设我们将它们分别拆分为 0.8、0.1、0.1 进行训练、测试和验证，您可以这样做这样:

train, test, val = np.split(df, [int(.8 * len(df)), int(.9 * len(df))])

我很想知道如何在使用这种方法分割数据时考虑分层。

Stratifying is splitting data while keeping the priors of each class you have in data. That is if you're going to take 0.8 for the training set, you take 0.8 from each class you have. Same for test and train.

我尝试首先按类别对数据进行分组:

grouped_df = df.groupby(class_col_name, group_keys=False)

但它没有显示正确的结果。

<小时/>

注意:我熟悉 train_test_split

Answer 1

只需使用您的 groupby 对象 grouped_df，它由每个子集数据帧组成，然后您可以在其中运行所需的 np.split。然后使用 pd.concat 连接所有采样数据帧。总之，这将根据您引用的消息分层:

train_list = []; test_list = [], val_list = []
grouped_df = df.groupby(class_col_name)

# ITERATE THROUGH EACH SUBSET DF
for i, g in grouped_df:
    # STRATIFY THE g (CLASS) DATA FRAME
    train, test, val = np.split(g, [int(.8 * len(g)), int(.9 * len(g))])

    train_list.append(train); test_list.append(test); val_list.append(val)

final_train = pd.concat(train_list)
final_test = pd.concat(test_list)
final_val = pd.concat(val_list)

或者，使用列表推导式的简写版本:

# LIST OF ARRAYS
arr_list = [np.split(g, [int(.8 * len(g)), int(.9 * len(g))]) for i, g in grouped_df]

final_train = pd.concat([t[0] for t in arr_list])
final_test = pd.concat([t[1] for t in arr_list])
final_val = pd.concat([v[2] for v in arr_list])