如何创建DataFrame存储训练-验证-测试集划分结果并生成id_set.csv

你现有的拆分逻辑已经符合比例要求：第一次拆分取20%作为测试集，剩下80%再拆分25%作为验证集，最终比例为训练集60%、验证集20%、测试集20%，正好落在你要求的区间内。
你可以按照以下步骤构造目标DataFrame：

1. 先确认原数据self.df中包含和dataset.csv完全一致的ID列，如果ID是DataFrame的索引，先执行代码转成普通列：

self.df = self.df.reset_index(names='ID')

2. 你原来的拆分代码不变，拆分完成后给三个子集分别打上set标签，再合并取需要的列导出即可：

# 原有拆分逻辑
train, test = train_test_split(self.df, test_size=0.2, random_state=1)
train, val = train_test_split(train, test_size=0.25, random_state=1)

# 给各子集添加set标识列
train['set'] = 'train'
val['set'] = 'validation'
test['set'] = 'test'

# 合并所有数据，只保留ID和set两列
id_set_df = pd.concat([train, val, test], axis=0)[['ID', 'set']]

# 导出为id_set.csv，不保留索引
id_set_df.to_csv('id_set.csv', index=False, encoding='utf-8')

导出的csv文件格式就完全符合你需要的结构，ID顺序不影响使用，如果需要按ID排序的话，可以在合并后加一行id_set_df = id_set_df.sort_values(by='ID', ignore_index=True)再导出。

内容的提问来源于stack exchange，提问作者melololo