如何按比例拆分未知列数的Pandas DataFrame为训练、测试、验证集

Hey there! 这个问题其实很好解决——不管你的CSV有3列还是40列，核心逻辑都是按行拆分样本，和列数完全无关。下面给你两种实用的实现方法，按需选择就行：

方法一：用sklearn的train_test_split（推荐） #

train_test_split是机器学习里常用的拆分工具，能帮你轻松实现分层拆分，还能保证随机性和可复现性。步骤很简单：先把数据集拆成「训练集（50%）」和「临时集（50%）」，再把临时集拆成「测试集（25%）」和「验证集（25%）」。

代码示例：

import pandas as pd
from sklearn.model_selection import train_test_split

# 读取CSV文件到DataFrame
df = pd.read_csv("your_dataset.csv")

# 第一步：拆分训练集和临时集（训练集占50%）
train_df, temp_df = train_test_split(df, test_size=0.5, random_state=42)

# 第二步：把临时集拆成测试集和验证集（各占原数据的25%）
test_df, val_df = train_test_split(temp_df, test_size=0.5, random_state=42)

关键说明： #

• random_state=42：固定随机种子，保证每次拆分结果完全一致，方便调试和复现；如果不需要固定结果，可以去掉这个参数。
• 如果你的数据集存在类别不平衡的情况，比如分类任务的标签分布不均，建议加上stratify参数，让每个子集的类别分布和原数据集一致：

  # 假设你的目标列叫'target'
  train_df, temp_df = train_test_split(df, test_size=0.5, random_state=42, stratify=df['target'])
  test_df, val_df = train_test_split(temp_df, test_size=0.5, random_state=42, stratify=temp_df['target'])
  

方法二：纯Pandas实现 #

如果你不想依赖sklearn，用Pandas自带的方法也能搞定：先打乱数据集，再按行数比例截取。

代码示例：

import pandas as pd

# 读取CSV文件到DataFrame
df = pd.read_csv("your_dataset.csv")

# 先打乱数据集（保证样本分布均匀）
shuffled_df = df.sample(frac=1, random_state=42).reset_index(drop=True)

# 计算各子集的行数
total_rows = shuffled_df.shape[0]
train_rows = int(total_rows * 0.5)
test_rows = int(total_rows * 0.25)

# 按行索引拆分
train_df = shuffled_df.iloc[:train_rows]
test_df = shuffled_df.iloc[train_rows:train_rows + test_rows]
val_df = shuffled_df.iloc[train_rows + test_rows:]

小提示： #

如果总行数不是4的倍数，拆分后比例会有微小偏差（比如少1行），如果追求绝对精确，可以用round()代替int()，不过一般这种微小误差对模型训练影响不大。

不管用哪种方法，所有列都会完整保留到训练/测试/验证集中，完全不用在意原数据集的列数多少～

内容的提问来源于stack exchange，提问作者João Victor Canabarro