我有一个大约200万条观测值的数据集，需要按照60:20:20的比例将其拆分为训练集、验证集和测试集。我的数据集的一个简化摘录如下：
+---------+------------+-----------+-----------+
| note_id | subject_id | category | note |
+---------+------------+-----------+-----------+
| 1 | 1 | ECG | blah ... |
| 2 | 1 | Discharge | blah ... |
| 3 | 1 | Nursing | blah ... |
| 4 | 2 | Nursing | blah ... |
| 5 | 2 | Nursing | blah ... |
| 6 | 3 | ECG | blah ... |
+---------+------------+-----------+-----------+数据集中有多个类别，这些类别分布不均衡，因此我需要确保训练集、验证集和测试集中各类别的比例与原始数据集相同。这部分没有问题，我可以使用sklearn库中的StratifiedShuffleSplit来实现。
然而，我还需要确保来自每个受试者的观测值不会在训练集、验证集和测试集中分开。来自某个特定受试者的所有观测值需要在同一个数据集中，以确保我的训练模型在验证/测试时从未见过该受试者。例如，每个subject_id 1的观测值都应该在训练集中。
我无法想出一种方法来确保按category进行分层拆分，防止subject_id在数据集之间“污染”（用词不当），确保60:20:20的拆分比例，并确保数据集以某种方式被打乱。任何帮助都将不胜感激！
谢谢！
编辑：
我现在了解到，通过sklearn库中的GroupShuffleSplit函数，也可以实现按类别分组并在数据集拆分时保持组的完整性。因此，本质上，我需要的是一种结合了分层和分组的打乱拆分，即StratifiedGroupShuffleSplit，但它并不存在。Github问题：https://github.com/scikit-learn/scikit-learn/issues/12076
回答：
本质上，我需要StratifiedGroupShuffleSplit，但它并不存在（Github问题）。这是因为这种功能的行为不明确，实现一个既分组又分层的拆分数据集并不总是可能的（也在此讨论）——尤其是像我这样严重不平衡的数据集。在我的情况下，我希望严格按照组进行分组，以确保组之间没有任何重叠，而分层和数据集的60:20:20比例拆分则可以尽可能近似地完成。
正如Ghanem提到的，我别无选择，只能自己构建一个函数来拆分数据集，我已经在下面完成了这个工作：
def StratifiedGroupShuffleSplit(df_main): df_main = df_main.reindex(np.random.permutation(df_main.index)) # 打乱数据集 # 创建空的训练集、验证集和测试集 df_train = pd.DataFrame() df_val = pd.DataFrame() df_test = pd.DataFrame() hparam_mse_wgt = 0.1 # 必须在0和1之间 assert(0 <= hparam_mse_wgt <= 1) train_proportion = 0.6 # 必须在0和1之间 assert(0 <= train_proportion <= 1) val_test_proportion = (1-train_proportion)/2 subject_grouped_df_main = df_main.groupby(['subject_id'], sort=False, as_index=False) category_grouped_df_main = df_main.groupby('category').count()[['subject_id']]/len(df_main)*100 def calc_mse_loss(df): grouped_df = df.groupby('category').count()[['subject_id']]/len(df)*100 df_temp = category_grouped_df_main.join(grouped_df, on = 'category', how = 'left', lsuffix = '_main') df_temp.fillna(0, inplace=True) df_temp['diff'] = (df_temp['subject_id_main'] - df_temp['subject_id'])**2 mse_loss = np.mean(df_temp['diff']) return mse_loss i = 0 for _, group in subject_grouped_df_main: if (i < 3): if (i == 0): df_train = df_train.append(pd.DataFrame(group), ignore_index=True) i += 1 continue elif (i == 1): df_val = df_val.append(pd.DataFrame(group), ignore_index=True) i += 1 continue else: df_test = df_test.append(pd.DataFrame(group), ignore_index=True) i += 1 continue mse_loss_diff_train = calc_mse_loss(df_train) - calc_mse_loss(df_train.append(pd.DataFrame(group), ignore_index=True)) mse_loss_diff_val = calc_mse_loss(df_val) - calc_mse_loss(df_val.append(pd.DataFrame(group), ignore_index=True)) mse_loss_diff_test = calc_mse_loss(df_test) - calc_mse_loss(df_test.append(pd.DataFrame(group), ignore_index=True)) total_records = len(df_train) + len(df_val) + len(df_test) len_diff_train = (train_proportion - (len(df_train)/total_records)) len_diff_val = (val_test_proportion - (len(df_val)/total_records)) len_diff_test = (val_test_proportion - (len(df_test)/total_records)) len_loss_diff_train = len_diff_train * abs(len_diff_train) len_loss_diff_val = len_diff_val * abs(len_diff_val) len_loss_diff_test = len_diff_test * abs(len_diff_test) loss_train = (hparam_mse_wgt * mse_loss_diff_train) + ((1-hparam_mse_wgt) * len_loss_diff_train) loss_val = (hparam_mse_wgt * mse_loss_diff_val) + ((1-hparam_mse_wgt) * len_loss_diff_val) loss_test = (hparam_mse_wgt * mse_loss_diff_test) + ((1-hparam_mse_wgt) * len_loss_diff_test) if (max(loss_train,loss_val,loss_test) == loss_train): df_train = df_train.append(pd.DataFrame(group), ignore_index=True) elif (max(loss_train,loss_val,loss_test) == loss_val): df_val = df_val.append(pd.DataFrame(group), ignore_index=True) else: df_test = df_test.append(pd.DataFrame(group), ignore_index=True) print ("Group " + str(i) + ". loss_train: " + str(loss_train) + " | " + "loss_val: " + str(loss_val) + " | " + "loss_test: " + str(loss_test) + " | ") i += 1 return df_train, df_val, df_testdf_train, df_val, df_test = StratifiedGroupShuffleSplit(df_main)我创建了一个基于以下两个方面的任意损失函数：
- 每个类别在百分比表示上的平均平方差与整体数据集的比较
- 数据集的比例长度与根据提供的比例（60:20:20）应有的长度之间的平方差
通过静态超参数hparam_mse_wgt来对损失函数的这两个输入进行加权。对于我的特定数据集，0.1的值效果很好，但我建议您在使用此函数时可以尝试调整。如果将其设置为0，将只关注保持拆分比例并忽略分层。如果将其设置为1，则反之亦然。
使用此损失函数，我然后遍历每个受试者（组），并根据哪个数据集的损失函数值最高，将其追加到相应的数据集（训练集、验证集或测试集）。
这并不特别复杂，但它对我来说完成了工作。它不一定适用于每个数据集，但数据集越大，成功的几率就越大。希望其他人会发现它有用。
