如何使用Scikit-learn拆分DataFrame为训练集与测试集

方法1：使用sklearn.model_selection.train_test_split（最常用方案） #

这是机器学习场景下的标准拆分方式，支持直接按比例分配，还能实现分层抽样来保证类别分布一致。

代码示例：

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设你的DataFrame名为df
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

• test_size=0.2：指定测试集占比20%，训练集自动为80%
• random_state=42：固定随机种子，确保每次拆分结果一致，方便复现实验
• 如果是分类任务需要分层抽样，可添加stratify=df['目标类别列名']参数，保证训练/测试集的类别比例和原数据一致

方法2：使用Pandas原生方法（无需依赖sklearn） #

如果不想引入额外依赖，用Pandas自带的采样和索引操作就能完成拆分：

import pandas as pd

# 按比例抽取80%作为训练集
train_df = df.sample(frac=0.8, random_state=42)
# 测试集取原数据中不在训练集里的样本
test_df = df.drop(train_df.index)

或者按样本数量计算拆分：

train_size = int(0.8 * len(df))
train_df = df.sample(n=train_size, random_state=42)
test_df = df.drop(train_df.index)

注意事项 #

• 针对大型DataFrame，两种方法都经过优化，不会过度占用内存
• 如果你的DataFrame存在重复索引，建议先执行df = df.reset_index(drop=True)，避免拆分时出现索引冲突
• 若数据是时间序列类型，不要用随机拆分，应该按时间顺序划分（比如取前80%时间范围的数据作为训练集）

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