我有一个训练数据框，每个班级都有相同数量的标签。这些类标有[-1,0,1,2,3]。我想构建一个验证集，其大小是训练集大小的20％，这样我在验证集中的每个类中都有相同数量的标签。我写了以下代码。 train_data_frame.stage具有标签[-1,0,1,2,3]：
total_stages_set = set(train_data_frame.stage)
minimum_samples = round(len(train_data_frame['stage'])/4 * 0.2)
validation_frames = [train_data_frame[train_data_frame.stage == s].sample(n=minimum_samples, replace=False) for s in total_stages_set]
validation_frame = pd.concat(validation_frames)
这给了我验证集。问题是我想从train_data_frame中删除这些行。我尝试使用行索引，但事实证明它们不是唯一的，因此删除了不需要的行。有什么建议？如果还有其他想法，我不介意废弃这段代码。