如何用Python将CSV文件按80/20拆分以用于机器学习训练

嘿，你已经掌握了逻辑回归、朴素贝叶斯和决策树的实现，就差数据拆分这一步啦！先说说你原来代码里的问题：你用read().split()读取文件后直接按前80个元素拆分，这其实是把文件内容按字节/空格拆成了零散的字符串片段，完全不是按样本（行）来划分的，而且没有随机化，很容易引入顺序偏差，肯定没法用在模型训练里。

下面给你两种靠谱的拆分方法，都是针对结构化CSV数据的正确姿势：

方法一：用scikit-learn的train_test_split（最推荐） #

这是机器学习领域的标准操作，它会帮你自动随机打乱数据，还支持分层抽样——比如你的糖尿病数据是分类任务，分层能保证训练集和测试集里的患病/不患病比例和原数据一致，避免拆分后的数据集类别失衡。

代码示例：

import pandas as pd
from sklearn.model_selection import train_test_split

# 读取CSV数据
diabetes_df = pd.read_csv("diabetes.csv")

# 先把特征和目标标签分开（假设最后一列是目标列，比如叫'Outcome'）
X = diabetes_df.drop('Outcome', axis=1)  # 所有特征列
y = diabetes_df['Outcome']  # 目标标签（比如是否患糖尿病）

# 执行拆分：80%训练集，20%测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,          # 测试集占比20%
    random_state=42,        # 随机种子，固定后每次拆分结果完全一致，方便复现
    stratify=y              # 分层抽样，保证类别比例和原数据一致
)

# 验证拆分结果
print(f"训练集样本数: {len(X_train)}，测试集样本数: {len(X_test)}")

方法二：手动用pandas拆分（适合理解原理） #

如果你想自己实现拆分逻辑，核心是先打乱数据，再按比例切割：

import pandas as pd

# 读取数据
diabetes_df = pd.read_csv("diabetes.csv")

# 第一步：随机打乱整个数据集（必须做，否则原数据的顺序会影响模型）
shuffled_df = diabetes_df.sample(frac=1, random_state=42)  # frac=1表示取全部数据，随机打乱

# 第二步：计算80%样本的位置
train_size = int(0.8 * len(shuffled_df))

# 第三步：拆分训练集和测试集
train_df = shuffled_df[:train_size]
test_df = shuffled_df[train_size:]

# 同样可以拆分特征和标签
X_train = train_df.drop('Outcome', axis=1)
y_train = train_df['Outcome']
X_test = test_df.drop('Outcome', axis=1)
y_test = test_df['Outcome']

关键注意点 #

• 一定要随机打乱数据：如果原数据是按类别排序的，直接按前80%拆分会导致训练集和测试集的类别分布不均，模型会失效。
• 分类任务尽量用分层抽样：比如你的糖尿病数据中，患病的样本可能占少数，分层能保证训练集和测试集都包含足够的正样本和负样本。

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