Skip to content

Sklearn train_test_split：Python中机器学习数据集划分

Soren Atelier
Name
Soren Atelier

Updated on

在用于评估的相同数据上训练机器学习模型会产生误导性的高准确率。模型记住了训练数据而不是学习可泛化的模式——这个问题被称为过拟合。你需要一个模型在训练过程中从未见过的独立测试集，以获得诚实的性能指标。

Scikit-learn的train_test_split()是将数据集划分为训练集和测试集的标准方法。它可以处理数组、DataFrame和稀疏矩阵，并提供分层抽样、可重复性和自定义划分比例的选项。

📚

    基本用法

    from sklearn.model_selection import train_test_split
import numpy as np
 
# 示例数据：100个样本，5个特征
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)  # 二分类标签
 
# 划分：80%训练集，20%测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
 
print(f"训练集：{X_train.shape[0]} 个样本")
print(f"测试集：{X_test.shape[0]} 个样本")
# 训练集：80 个样本
# 测试集：20 个样本

    关键参数

    参数默认值描述
    test_size0.25测试样本的比例(0.0-1.0)或绝对数量
    train_sizeNone训练样本的比例或数量(test_size的补集)
    random_stateNone可重复划分的种子
    shuffleTrue划分前是否打乱数据
    stratifyNone用于分层划分的数组

    与Pandas DataFrame一起使用

    from sklearn.model_selection import train_test_split
import pandas as pd
 
df = pd.DataFrame({
    'age': [25, 30, 35, 40, 45, 50, 55, 60, 28, 33],
    'income': [40, 50, 60, 70, 80, 90, 100, 110, 45, 55],
    'purchased': [0, 0, 1, 1, 1, 1, 1, 1, 0, 0]
})
 
X = df[['age', 'income']]
y = df['purchased']
 
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)
 
print(f"训练集形状：{X_train.shape}")  # (7, 2)
print(f"测试集形状：{X_test.shape}")   # (3, 2)
print(type(X_train))  # <class 'pandas.core.frame.DataFrame'>

    划分后DataFrame仍保持为DataFrame——列名和索引被保留。

    random_state：可重复的划分

    没有random_state时，每次运行会得到不同的划分：

    from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.arange(10).reshape(5, 2)
y = np.array([0, 0, 1, 1, 1])
 
# 不设random_state：每次运行不同的划分
_, X_test1, _, _ = train_test_split(X, y, test_size=0.4)
_, X_test2, _, _ = train_test_split(X, y, test_size=0.4)
print(np.array_equal(X_test1, X_test2))  # 很可能是False
 
# 设置random_state：每次相同的划分
_, X_test3, _, _ = train_test_split(X, y, test_size=0.4, random_state=42)
_, X_test4, _, _ = train_test_split(X, y, test_size=0.4, random_state=42)
print(np.array_equal(X_test3, X_test4))  # True

    **始终设置random_state以确保可重复性。**使用任何整数——42是惯例，但具体数字并不重要。

    分层划分

    对于不平衡数据集，随机划分可能会将大多数少数类样本放在一个集合中。分层确保两个集合具有相同的类别比例：

    from sklearn.model_selection import train_test_split
import numpy as np
from collections import Counter
 
# 不平衡数据集：90%类别0，10%类别1
np.random.seed(42)
X = np.random.randn(200, 4)
y = np.array([0] * 180 + [1] * 20)
 
# 不使用分层
_, _, y_train_bad, y_test_bad = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print("不使用分层：")
print(f"  训练集：{Counter(y_train_bad)}")
print(f"  测试集：{Counter(y_test_bad)}")
 
# 使用分层
_, _, y_train_good, y_test_good = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print("\n使用stratify=y：")
print(f"  训练集：{Counter(y_train_good)}")
print(f"  测试集：{Counter(y_test_good)}")
# 训练集：Counter({0: 144, 1: 16})  -- 10%类别1
# 测试集：Counter({0: 36, 1: 4})    -- 10%类别1

    何时使用分层

    场景使用分层？
    平衡类别 (50/50)可选
    不平衡类别 (90/10)是
    多类别分类是
    回归（连续目标）否（不支持）
    小数据集（< 100个样本）是（防止空类别）

    训练集/验证集/测试集划分

    对于超参数调优，你需要三个集合：训练集、验证集和测试集。调用train_test_split两次：

    from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.random.randn(1000, 10)
y = np.random.randint(0, 3, 1000)
 
# 第一次划分：80%训练集+验证集，20%测试集
X_temp, X_test, y_temp, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# 第二次划分：75%训练集，25%验证集（占80%的25% = 总体60/20/20）
X_train, X_val, y_train, y_val = train_test_split(
    X_temp, y_temp, test_size=0.25, random_state=42, stratify=y_temp
)
 
print(f"训练集：    {X_train.shape[0]} 个样本 (60%)")
print(f"验证集：    {X_val.shape[0]} 个样本 (20%)")
print(f"测试集：    {X_test.shape[0]} 个样本 (20%)")

    完整的机器学习管道示例

    from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
import numpy as np
 
# 生成示例数据
np.random.seed(42)
X = np.random.randn(500, 8)
y = (X[:, 0] + X[:, 1] > 0).astype(int)
 
# 划分数据
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# 特征缩放（只在训练集上fit！）
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)  # 使用训练集的统计量
 
# 训练模型
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
 
# 评估
y_pred = model.predict(X_test_scaled)
print(f"准确率：{accuracy_score(y_test, y_pred):.4f}")
print(classification_report(y_test, y_pred))

    **关键规则：**只在训练集上拟合缩放器（以及任何预处理）。将相同的变换应用于测试集。在完整数据集上拟合会导致数据泄漏。

    常用划分比例

    划分训练测试何时使用
    80/2080%20%默认选择，大多数数据集
    70/3070%30%小数据集，需要更大的测试集
    90/1090%10%大数据集（10k+样本）
    60/20/2060%20%验证 + 20%测试调优超参数时

    探索模型结果

    划分和训练模型之后，PyGWalker (opens in a new tab)让你可以在Jupyter中交互式地探索预测与实际值、训练/测试集的特征分布以及错误模式：

    import pandas as pd
import pygwalker as pyg
 
results = pd.DataFrame({
    'actual': y_test,
    'predicted': y_pred,
    'correct': y_test == y_pred
})
walker = pyg.walk(results)

    常见问题

    sklearn中的train_test_split做什么？

    train_test_split()将数组或DataFrame随机分为两个子集：一个用于训练，一个用于测试。它确保模型评估使用模型在训练期间从未见过的数据，从而给出诚实的性能估计。

    最佳的训练/测试划分比例是多少？

    80/20是标准默认值。对于需要可靠测试集的较小数据集使用70/30，对于大数据集（10k+样本）使用90/10，其中10%仍然相当可观。对于超参数调优，使用60/20/20（训练/验证/测试）。

    train_test_split中的random_state有什么作用？

    random_state设置划分前打乱的随机种子。使用相同的random_state值每次都会产生相同的划分，使结果可重复。任何整数都可以。

    什么时候应该在train_test_split中使用stratify？

    当目标变量不平衡（例如95%负样本，5%正样本）或数据集较小时使用stratify=y。分层确保训练集和测试集具有相同比例的每个类别。

    如何将数据划分为训练集、验证集和测试集？

    调用train_test_split两次。首先划分为训练集+验证集和测试集（例如80/20）。然后将训练集+验证集划分为训练集和验证集（例如80%的75/25，总体为60/20/20）。或者，使用sklearn.model_selection.KFold进行交叉验证。

    总结

    train_test_split()是Python中每个机器学习工作流的基础。始终在训练前使用它——永远不要在训练数据上评估。设置random_state以确保可重复性，对不平衡类别使用stratify=y，并记住只在训练集上拟合预处理步骤。对于模型选择，三方划分（训练/验证/测试）或使用交叉验证以获得更稳健的估计。

    📚