Skip to content
话题
Python
Sklearn Train Test Split:Python数据分割完整指南

Sklearn Train Test Split:Python数据分割完整指南

Soren Atelier
Name
Soren Atelier

Updated on

在整个数据集上训练机器学习模型然后在相同数据上评估会导致一个关键问题:你的模型看起来表现良好,但实际上只是记住了数据而不是学习模式。这种过拟合意味着当模型遇到新的、未见过的数据时会失败得很惨。数据科学家需要一种可靠的方法来评估模型在训练期间从未见过的数据上的性能。

解决方案是训练-测试分割。通过保留一部分数据用于评估,你可以获得模型在真实世界中如何表现的诚实评估。sklearn的train_test_split函数使这个过程变得简单,但不正确使用仍然可能导致数据泄漏、泛化能力差和误导性的性能指标。

本指南涵盖了关于sklearn的train_test_split你需要知道的一切,从基本用法到时间序列数据、不平衡类和多输出问题的高级技术。

📚

    什么是Train Test Split?

    Train test split是评估机器学习模型的基本技术。你将数据集分成两部分:用于拟合模型的训练集,以及用于评估模型在未见数据上性能的测试集。

    scikit-learn(sklearn)的train_test_split函数自动化了这个过程,只用一行代码就能处理随机打乱和分割。

    from sklearn.model_selection import train_test_split
 
# 基本用法
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    在这个例子中,X包含你的特征(输入变量),y包含你的目标变量(你想要预测的内容)。函数返回四个数组:训练特征、测试特征、训练标签和测试标签。

    基本train_test_split语法

    train_test_split的最简单用法只需要两个参数:你的特征和目标变量。

    import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
 
# 加载示例数据
iris = load_iris()
X = iris.data
y = iris.target
 
# 分割数据(80%训练,20%测试)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
 
print(f"Training samples: {len(X_train)}")
print(f"Test samples: {len(X_test)}")

    这会随机分割你的数据,80%用于训练,20%用于测试。然而,这种基本用法有一个关键缺陷:每次运行代码时分割都不同,使结果无法重现。

    基本参数

    test_size和train_size

    test_size参数控制有多少数据进入测试集。你可以将其指定为:

    • 0.0到1.0之间的浮点数(数据集的比例)
    • 整数(测试样本的绝对数量)
    # 30%测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)
 
# 测试集中50个样本
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=50)
 
# 或者,指定train_size
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8)

    如果同时指定test_size和train_size,它们必须加起来等于1.0(或如果使用整数则等于总数据集大小)。在大多数情况下,只指定test_size就足够了。

    random_state用于可重现性

    random_state参数对于可重现的结果至关重要。没有它,每次运行代码都会得到不同的分割,使得调试或比较实验变得不可能。

    # 没有random_state - 每次分割都不同
X_train1, X_test1, y_train1, y_test1 = train_test_split(X, y, test_size=0.2)
X_train2, X_test2, y_train2, y_test2 = train_test_split(X, y, test_size=0.2)
 
print(f"Same split? {(X_train1 == X_train2).all()}")  # False
 
# 有random_state - 每次分割都相同
X_train1, X_test1, y_train1, y_test1 = train_test_split(X, y, test_size=0.2, random_state=42)
X_train2, X_test2, y_train2, y_test2 = train_test_split(X, y, test_size=0.2, random_state=42)
 
print(f"Same split? {(X_train1 == X_train2).all()}")  # True

    为random_state使用任何整数。具体数字并不重要;重要的是在整个项目中一致地使用相同的数字。

    shuffle参数

    默认情况下,train_test_split在分割前会打乱数据。对于大多数机器学习任务,这正是你想要的。但是,对于时间序列数据或当顺序很重要时,你应该禁用打乱。

    # 启用打乱(默认)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=True)
 
# 禁用打乱(用于时间序列)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

    当shuffle=False时,函数简单地将前面部分用于训练,后面部分用于测试,保持原始顺序。

    参数参考表

    参数类型默认值描述
    test_sizefloat或intNone测试集的比例(0.0-1.0)或样本数
    train_sizefloat或intNone训练集的比例(0.0-1.0)或样本数
    random_stateintNone用于可重现性的随机种子
    shuffleboolTrue分割前是否打乱数据
    stratifyarray-likeNone用于分层分割的数据

    不平衡数据的分层分割

    当数据集有不平衡的类别(某些类别的样本远少于其他类别)时,随机分割可能创建训练或测试集,这些集合不能很好地代表整体分布。这对分类任务尤其有问题。

    stratify参数确保训练集和测试集中的类别分布与原始数据集匹配。

    import numpy as np
from sklearn.model_selection import train_test_split
 
# 创建不平衡数据集(90%类别0,10%类别1)
X = np.random.randn(1000, 5)
y = np.array([0] * 900 + [1] * 100)
 
# 不分层
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print(f"Train distribution: Class 0: {sum(y_train == 0)}, Class 1: {sum(y_train == 1)}")
print(f"Test distribution: Class 0: {sum(y_test == 0)}, Class 1: {sum(y_test == 1)}")
 
# 分层
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\nStratified train distribution: Class 0: {sum(y_train == 0)}, Class 1: {sum(y_train == 1)}")
print(f"Stratified test distribution: Class 0: {sum(y_test == 0)}, Class 1: {sum(y_test == 1)}")

    使用分层,训练集和测试集都保持90/10的类别分布。不使用分层,你可能运气好得到有代表性的分割,或者可能最终得到只有5%类别1的测试集,导致不可靠的评估指标。

    分割多个数组

    你可以一次分割多个数组,sklearn会确保它们以相同的方式分割(所有数组使用相同的索引)。

    import numpy as np
 
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)
sample_weights = np.random.rand(100)
 
# 分割所有三个数组
X_train, X_test, y_train, y_test, weights_train, weights_test = train_test_split(
    X, y, sample_weights, test_size=0.2, random_state=42
)
 
print(f"X_train shape: {X_train.shape}")
print(f"y_train shape: {y_train.shape}")
print(f"weights_train shape: {weights_train.shape}")

    这在你有样本权重、多个目标变量或需要一致分割的附加元数据时特别有用。

    Train/Test Split vs 交叉验证 vs Holdout

    不同的验证策略服务于不同的目的。这里是比较:

    方法数据使用计算成本最适合限制
    Train/Test Split70-80%训练,20-30%测试低快速模型评估,大型数据集单次评估,分割可能幸运/不幸
    交叉验证100%用于训练/测试(k-fold)高(k倍慢)小型数据集,可靠的性能估计计算成本高,不适合时间序列
    Train/Val/Test (Holdout)60%训练,20%验证,20%测试中等超参数调优,最终评估需要更多数据,工作流更复杂
    from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
 
iris = load_iris()
X, y = iris.data, iris.target
model = RandomForestClassifier(random_state=42)
 
# 方法1:简单的train/test分割
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model.fit(X_train, y_train)
print(f"Train/Test Split Score: {model.score(X_test, y_test):.3f}")
 
# 方法2:5折交叉验证
scores = cross_val_score(model, X, y, cv=5)
print(f"Cross-Validation Score: {scores.mean():.3f} (+/- {scores.std():.3f})")
 
# 方法3:Train/Val/Test分割
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42)
model.fit(X_train, y_train)
print(f"Validation Score: {model.score(X_val, y_val):.3f}")
print(f"Test Score: {model.score(X_test, y_test):.3f}")

    对于大多数项目,从简单的train/test分割开始。当数据有限或需要更稳健的性能估计时使用交叉验证。当需要调整超参数时使用train/val/test。

    高级分割技术

    时间序列分割

    对于时间序列数据,随机打乱会破坏时间顺序,可能导致数据泄漏(使用未来信息预测过去)。请改用TimeSeriesSplit:

    from sklearn.model_selection import TimeSeriesSplit
import numpy as np
 
X = np.random.randn(100, 5)
y = np.random.randn(100)
 
tscv = TimeSeriesSplit(n_splits=5)
 
for train_index, test_index in tscv.split(X):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
    print(f"Train size: {len(train_index)}, Test size: {len(test_index)}")

    TimeSeriesSplit创建多个train/test分割,其中每个训练集包括到某个时间点的所有过去数据,测试集包括紧随其后的时期。这模拟了现实世界的预测,其中你只有过去的数据来预测未来。

    分组数据的GroupShuffleSplit

    当你的数据有分组(例如,来自同一患者的多次测量、来自同一客户的多次交易)时,你需要确保整个组保持在训练集或测试集中,以避免数据泄漏。

    from sklearn.model_selection import GroupShuffleSplit
import numpy as np
 
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)
groups = np.array([0] * 25 + [1] * 25 + [2] * 25 + [3] * 25)  # 4个组
 
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
 
for train_idx, test_idx in gss.split(X, y, groups):
    X_train, X_test = X[train_idx], X[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]
    print(f"Train groups: {np.unique(groups[train_idx])}")
    print(f"Test groups: {np.unique(groups[test_idx])}")

    这确保给定组的所有样本都在训练集或测试集中,永远不会同时在两者中。

    Train Test Split最佳实践

    选择正确的分割比例

    最常见的分割比例是:

    • 80/20:中大型数据集(10,000+样本)的标准选择
    • 70/30:对于较小的数据集(1,000-10,000样本)更好,以获得更稳健的测试评估
    • 90/10:对于非常大的数据集(100,000+样本),其中即使10%也提供充足的测试样本
    • 60/20/20:调整超参数时用于train/validation/test
    import numpy as np
 
def recommend_split_ratio(n_samples):
    if n_samples < 1000:
        return "Consider cross-validation instead of simple split"
    elif n_samples < 10000:
        return "70/30 split recommended"
    elif n_samples < 100000:
        return "80/20 split recommended"
    else:
        return "90/10 or 80/20 split recommended"
 
sample_sizes = [500, 5000, 50000, 500000]
for size in sample_sizes:
    print(f"{size} samples: {recommend_split_ratio(size)}")

    避免数据泄漏

    当测试集的信息影响训练过程时会发生数据泄漏。常见来源:

    1. 分割前预处理:始终先分割,然后预处理
    2. 对组合数据进行特征缩放:只在训练数据上拟合缩放器
    3. 对组合数据进行特征选择:只使用训练数据选择特征
    from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.random.randn(1000, 10)
y = np.random.randint(0, 2, 1000)
 
# 错误:分割前缩放(数据泄漏!)
scaler_wrong = StandardScaler()
X_scaled_wrong = scaler_wrong.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled_wrong, y, test_size=0.2)
 
# 正确:先分割,然后缩放
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)  # 在训练数据上拟合
X_test_scaled = scaler.transform(X_test)  # 使用训练统计信息转换测试数据

    错误的方法使用整个数据集(包括测试样本)的信息来计算缩放参数,这会将测试集的信息泄漏到训练过程中。

    尽可能分层

    对于分类问题,除非有特定原因,否则始终使用分层分割。这对以下情况特别重要:

    • 不平衡的数据集
    • 小型数据集
    • 有罕见类别的多类问题
    from sklearn.model_selection import train_test_split
import numpy as np
 
# 罕见疾病数据集:1%阳性病例
X = np.random.randn(1000, 20)
y = np.array([0] * 990 + [1] * 10)
 
# 不分层 - 测试集中可能没有阳性病例!
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=123)
print(f"Non-stratified test positives: {sum(y_test)}")
 
# 分层 - 保证比例表示
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.1, random_state=123, stratify=y
)
print(f"Stratified test positives: {sum(y_test)}")

    要避免的常见错误

    1. 忘记random_state

    没有random_state,每次运行代码时结果都会改变。这使得调试变得不可能,实验无法重现。

    # 差:没有random_state
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
 
# 好:设置random_state
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    2. 不对不平衡类别分层

    对于不平衡的数据集,随机分割可能创建高度不具代表性的测试集,导致不可靠的性能指标。

    # 差:不平衡数据没有分层
X_train, X_test, y_train, y_test = train_test_split(X, y_imbalanced, test_size=0.2)
 
# 好:使用分层
X_train, X_test, y_train, y_test = train_test_split(
    X, y_imbalanced, test_size=0.2, stratify=y_imbalanced, random_state=42
)

    3. 用Shuffle分割时间序列数据

    时间序列模型依赖于时间顺序。打乱会破坏这种结构,可能导致严重的数据泄漏。

    # 差:打乱时间序列数据
X_train, X_test, y_train, y_test = train_test_split(
    X_timeseries, y_timeseries, test_size=0.2, shuffle=True
)
 
# 好:禁用打乱或使用TimeSeriesSplit
X_train, X_test, y_train, y_test = train_test_split(
    X_timeseries, y_timeseries, test_size=0.2, shuffle=False
)

    4. 分割前预处理

    在分割前对整个数据集拟合预处理器(缩放器、填充器、编码器)会导致数据泄漏。

    # 差:分割前预处理
X_scaled = StandardScaler().fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2)
 
# 好:先分割,然后预处理
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

    5. 使用测试集调整超参数

    测试集应该只用于最终评估。如果你用它来选择超参数,本质上就是在测试数据上训练。

    # 差:在测试集上调优
from sklearn.ensemble import RandomForestClassifier
 
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
best_score = 0
best_params = None
for n_estimators in [10, 50, 100]:
    model = RandomForestClassifier(n_estimators=n_estimators)
    model.fit(X_train, y_train)
    score = model.score(X_test, y_test)  # 使用测试集!
    if score > best_score:
        best_score = score
        best_params = n_estimators
 
# 好:使用验证集或交叉验证
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42)
 
best_score = 0
best_params = None
for n_estimators in [10, 50, 100]:
    model = RandomForestClassifier(n_estimators=n_estimators)
    model.fit(X_train, y_train)
    score = model.score(X_val, y_val)  # 使用验证集
    if score > best_score:
        best_score = score
        best_params = n_estimators

    实用示例:完整工作流程

    这是正确使用train_test_split的完整机器学习工作流程:

    import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
 
# 加载数据
np.random.seed(42)
X = np.random.randn(1000, 10)
y = (X[:, 0] + X[:, 1] > 0).astype(int)  # 二分类
 
# 步骤1:分割数据(分层以获得平衡的测试集)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# 步骤2:预处理(只在训练数据上拟合)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
 
# 步骤3:训练模型
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
 
# 步骤4:评估
y_pred = model.predict(X_test_scaled)
print("Classification Report:")
print(classification_report(y_test, y_pred))
print("\nConfusion Matrix:")
print(confusion_matrix(y_test, y_pred))
 
# 步骤5:检查过拟合
train_score = model.score(X_train_scaled, y_train)
test_score = model.score(X_test_scaled, y_test)
print(f"\nTrain accuracy: {train_score:.3f}")
print(f"Test accuracy: {test_score:.3f}")
print(f"Overfitting gap: {train_score - test_score:.3f}")

    使用RunCell进行交互式数据分割

    在Jupyter笔记本中工作时,尝试不同的分割比例和参数可能很繁琐。RunCell (opens in a new tab)提供了专门为Jupyter中的数据科学工作流程设计的AI代理。它可以帮助你:

    • 自动测试多个分割比例并比较结果
    • 检测预处理管道中的数据泄漏
    • 为特定数据集建议最佳分层策略
    • 生成验证曲线以选择正确的train/test比例

    RunCell直接集成到你的Jupyter环境中,使你可以轻松迭代数据分割策略,而无需编写重复代码。

    使用PyGWalker可视化数据

    分割数据后,验证训练集和测试集具有相似的分布至关重要。PyGWalker (opens in a new tab)将你的pandas DataFrames转换为Tableau风格的交互式可视化,使以下操作变得容易:

    • 比较训练集和测试集之间的特征分布
    • 识别分割中的潜在采样偏差
    • 可视化类别不平衡并验证分层是否正确工作
    • 探索训练数据中特征之间的关系
    import pygwalker as pyg
import pandas as pd
 
# 转换为DataFrames用于可视化
train_df = pd.DataFrame(X_train, columns=[f'feature_{i}' for i in range(X_train.shape[1])])
train_df['dataset'] = 'train'
test_df = pd.DataFrame(X_test, columns=[f'feature_{i}' for i in range(X_test.shape[1])])
test_df['dataset'] = 'test'
 
combined = pd.concat([train_df, test_df])
 
# 创建交互式可视化
pyg.walk(combined)

    这让你可以交互式地探索训练和测试分布是否匹配,这对可靠的模型评估至关重要。

    常见问题

    如何在80/20和70/30分割之间选择?

    对于超过10,000个样本的数据集使用80/20,对于较小的数据集(1,000-10,000个样本)使用70/30。关键是确保测试集有足够的样本用于可靠评估——通常分类问题至少需要200-500个样本。对于非常大的数据集(100,000+样本),你可以使用90/10甚至95/5,因为即使5%也能提供数千个测试样本。

    random_state是什么,为什么重要?

    random_state是在分割前打乱数据的随机数生成器的种子。使用相同的random_state值确保每次运行代码时获得相同的分割,这对于可重现性和调试至关重要。没有它,每次都会得到不同的train/test分割,使得无法确定性能变化是由于模型改进还是只是幸运/不幸的数据分割。

    何时应该使用stratify参数?

    对所有分类问题使用stratify=y,特别是当你有不平衡的类别或小型数据集时。分层确保训练集和测试集中的类别分布与整体分布匹配。例如,如果你的数据中10%是阳性病例,分层保证训练集和测试集都有大约10%的阳性病例,防止由于不具代表性的分割而产生的评估偏差。

    可以对时间序列数据使用train_test_split吗?

    不,你不应该对时间序列数据使用带有shuffle=True的train_test_split,因为它会破坏时间顺序并导致数据泄漏(使用未来数据预测过去)。相反,对于简单的时间顺序分割,使用带有shuffle=False的train_test_split,或者使用TimeSeriesSplit进行尊重时间顺序的交叉验证。对于时间序列,始终确保训练数据在时间上先于测试数据。

    train_test_split与交叉验证有何不同?

    train_test_split创建单个train/test分区(通常为80/20),给你一个性能估计。交叉验证(如k-fold)创建多个train/test分割并平均结果,提供更稳健的性能估计。对于快速评估和大型数据集使用train_test_split。对于小型数据集(少于1,000个样本)或需要更可靠的性能估计时使用交叉验证。交叉验证慢k倍(例如,5-fold慢5倍)但减少性能指标的方差。

    结论

    Sklearn的train_test_split是评估机器学习模型的基础工具。通过正确分割数据,你可以获得预测真实世界模型行为的诚实性能估计。记住关键原则:始终设置random_state以确保可重现性,对分类问题使用stratify,避免分割前预处理,并根据数据集大小选择分割比例。

    掌握这些基础知识,你将避免导致过拟合模型和误导性性能指标的最常见陷阱。无论你是构建简单的分类器还是复杂的深度学习系统,适当的train-test分割都是可靠机器学习的第一步。

    📚
      Sklearn Train Test Split: Complete Guide to Splitting Data in PythonSnowflake Connector Python: Install and Connect to Snowflake with Ease