Mirko Stojiljković
通过 scikit-learn 中的 train_test_split() 函数,你可以高效地将数据集划分为训练集和测试集,从而确保在机器学习中对模型进行无偏评估。这种划分方式有助于防止过拟合和欠拟合,因为测试数据与训练数据是相互隔离的,使你能够准确评估模型的预测性能。
完成本教程后,你将理解以下内容:
train_test_split()是 sklearn 中用于将数据集划分为训练集和测试集的函数。x_train和y_train分别表示训练数据子集的输入和输出,而x_test和y_test则分别表示测试数据子集的输入和输出。- 通过指定
test_size=0.2,你可以使用数据集的 20% 作为测试集,剩余 80% 用于训练。 train_test_split()可以通过stratify参数处理不平衡数据集,以保持各类别的分布比例。- 你将学会如何使用
train_test_split()并将这些概念应用于现实场景,确保你的机器学习模型得到精确且公平的评估。此外,你还将探索sklearn.model_selection中的相关工具以获得更深入的理解。
数据划分的重要性
监督式机器学习的目标是创建能够将给定输入精确映射到给定输出的模型。输入也称为自变量或预测变量,而输出则可称为因变量或响应变量。
衡量模型精度的方式取决于你要解决的问题类型。在回归分析中,通常使用决定系数(R²)、均方根误差(RMSE)、平均绝对误差(MAE)或类似指标。对于分类问题,则常使用准确率(accuracy)、精确率(precision)、召回率(recall)、F1 分数及相关指标。
不同领域对这些精度度量指标的可接受数值范围各不相同。你可以在 Statistics By Jim、Quora 以及许多其他资源中找到详细解释。
最重要的是要理解:通常你需要无偏评估才能正确使用这些指标、评估模型的预测性能并验证模型。
这意味着你不能用训练模型时使用的相同数据来评估模型的预测性能。你需要使用模型从未见过的新数据来进行评估。你可以通过在使用数据之前先将其拆分来实现这一点。
训练集、验证集和测试集
为了对预测性能进行无偏评估,拆分数据集至关重要。在大多数情况下,将数据集随机划分为三个子集就足够了:
- 训练集:用于训练或拟合你的模型。例如,你可以使用训练集为线性回归、逻辑回归或神经网络找到最优权重(或系数)。
- 验证集:用于在超参数调优过程中对模型进行无偏评估。例如,当你想确定神经网络中的最优神经元数量或支持向量机的最佳核函数时,你会尝试不同的设置。对于每组超参数配置,你都使用训练集拟合模型,并用验证集评估其性能。
- 测试集:用于对最终模型进行无偏评估。你不应该将其用于拟合或验证。
在不太复杂的情况下(即不需要调优超参数时),仅使用训练集和测试集即可。
欠拟合与过拟合
拆分数据集对于检测模型是否出现两种非常常见的问题——欠拟合(underfitting)和过拟合(overfitting)——也很重要:
- 欠拟合通常是由于模型无法捕捉数据之间的关系所致。例如,试图用线性模型表示非线性关系时就会发生这种情况。欠拟合的模型在训练集和测试集上的表现通常都很差。
- 过拟合通常发生在模型结构过于复杂时,它不仅学到了数据中真实的关系,还学到了噪声。这类模型的泛化能力往往较差。虽然它们在训练数据上表现良好,但在未见过的测试数据上通常表现不佳。
你可以在《Python 中的线性回归》一文中找到关于欠拟合和过拟合的更详细解释。
使用 train_test_split() 的前提条件
现在你已经理解了拆分数据集的必要性——用于无偏模型评估以及识别欠拟合或过拟合问题——接下来就可以学习如何拆分自己的数据集了。
你将使用 scikit-learn 1.5.0 版本(简称 sklearn)。该库包含许多用于数据科学和机器学习的包,但本教程将重点关注 model_selection 包,特别是其中的 train_test_split() 函数。
注意:虽然本教程是在特定版本的 scikit-learn 上测试的,但你将使用的功能是该库的核心功能,在其他版本的 scikit-learn 中也应该能正常工作。
你可以使用 pip 安装 sklearn:
$ python -m pip install "scikit-learn==1.5.0"
如果你使用 Anaconda,可能已经安装了它。但如果你想使用一个干净的环境,请确保你拥有指定版本,或者使用 Miniconda。然后你可以从 Anaconda Cloud 使用 conda 安装:
$ conda install -c anaconda scikit-learn=1.5.0
你还需要 NumPy,但无需单独安装。如果你尚未安装 NumPy,安装 sklearn 时会自动安装。如果你想复习 NumPy 知识,可以参考《NumPy 教程:你进入 Python 数据科学的第一步》。
train_test_split() 的应用
在使用之前,你需要先导入 train_test_split() 和 NumPy。你可以在 Jupyter Notebook 中操作,也可以启动一个新的 Python REPL 会话,然后开始导入:
>>> import numpy as np
>>> from sklearn.model_selection import train_test_split
现在你已经导入了所需模块,就可以使用它们将数据划分为训练集和测试集了。你可以在一次函数调用中同时拆分输入和输出。
使用 train_test_split() 时,你只需提供想要拆分的数组。此外,你还可以提供一些可选参数。该函数通常返回一个 NumPy 数组列表,但在适当情况下也可以返回其他可迭代类型,例如 SciPy 稀疏矩阵:
sklearn.model_selection.train_test_split(*arrays, **options) -> list
函数签名中的 arrays 参数指的是包含你要拆分数据的列表、NumPy 数组、pandas DataFrame 或其他类似数组的对象序列。所有这些对象共同构成数据集,且长度必须相同。
在监督式机器学习应用中,你通常会处理两个这样的数组:
- 一个二维数组,包含输入(x)
- 一个一维数组,包含输出(y)
options 参数表示你可以使用可选关键字参数来自定义函数的行为:
train_size:定义训练集的大小。如果提供浮点数,则必须在 0.0 到 1.0 之间,表示用于训练的数据集比例;如果提供整数,则表示训练样本的总数。默认值为None。test_size:定义测试集的大小。与train_size非常相似。你应该提供train_size或test_size中的一个。如果两者都未提供,则默认使用数据集的 25%(即 0.25)用于测试。random_state:控制拆分过程中随机化的对象。可以是整数或RandomState实例。如果你需要结果可复现,则应设置随机状态。默认值为None。shuffle:布尔值,决定是否在拆分前打乱数据集。默认值为True。stratify:一个类数组对象,如果不为None,则决定如何进行分层拆分。
现在是时候尝试数据拆分了!你将从创建一个简单的数据集开始。该数据集将在二维数组 x 中包含输入,在一维数组 y 中包含输出:
>>> x = np.arange(1, 25).reshape(12, 2)
>>> y = np.array([0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 0])
>>> x
array([[ 1, 2],
[ 3, 4],
[ 5, 6],
[ 7, 8],
[ 9, 10],
[11, 12],
[13, 14],
[15, 16],
[17, 18],
[19, 20],
[21, 22],
[23, 24]])
>>> y
array([0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 0])
为了生成数据,你使用了 arange(),它非常适合基于数值范围生成数组。你还使用了 .reshape() 来修改 arange() 返回数组的形状,从而获得二维数据结构。
你可以在一次函数调用中同时拆分输入和输出数据集:
>>> x_train, x_test, y_train, y_test = train_test_split(x, y)
>>> x_train
array([[15, 16],
[21, 22],
[11, 12],
[17, 18],
[13, 14],
[ 9, 10],
[ 1, 2],
[ 3, 4],
[19, 20]])
>>> x_test
array([[ 5, 6],
[ 7, 8],
[23, 24]])
>>> y_train
array([1, 1, 0, 1, 0, 1, 0, 1, 0])
>>> y_test
array([1, 0, 0])
给定两个数组(如这里的 x 和 y),train_test_split() 会执行拆分并按以下顺序返回四个数组(在此例中为 NumPy 数组):
x_train:第一个数组(x)的训练部分x_test:第一个数组(x)的测试部分y_train:第二个数组(y)的训练部分y_test:第二个数组(y)的测试部分
你得到的结果可能与这里显示的不同。这是因为数据集拆分默认是随机的,每次运行函数时结果都会不同。然而,这通常不是你想要的。
有时,为了使测试结果可复现,你需要每次函数调用都得到相同的随机拆分结果。你可以通过 random_state 参数实现这一点。random_state 的具体值并不重要——它可以是任何非负整数。你也可以使用 numpy.random.RandomState 的实例,但那是更复杂的方法。
在前面的例子中,你使用了一个包含 12 行(或观测值)的数据集,并得到了一个包含 9 行的训练样本和一个包含 3 行的测试样本。这是因为你没有指定训练集和测试集的期望大小。默认情况下,25% 的样本被分配给测试集。这个比例在许多应用中通常是合适的,但并不总是符合你的需求。
通常,你会希望显式定义测试集或训练集的大小,有时甚至想尝试不同的值。你可以使用 train_size 或 test_size 参数来实现这一点。
修改代码,以便你可以选择测试集的大小并获得可复现的结果:
>>> x_train, x_test, y_train, y_test = train_test_split(
... x, y, test_size=4, random_state=4
... )
>>> x_train
array([[17, 18],
[ 5, 6],
[23, 24],
[ 1, 2],
[ 3, 4],
[11, 12],
[15, 16],
[21, 22]])
>>> x_test
array([[ 7, 8],
[ 9, 10],
[13, 14],
[19, 20]])
>>> y_train
array([1, 1, 0, 0, 1, 0, 1, 1])
>>> y_test
array([0, 1, 0, 0])
通过这一更改,你得到了与之前不同的结果。之前你有 9 个训练项和 3 个测试项。现在,由于参数 test_size=4,训练集有 8 个项目,测试集有 4 个项目。如果你使用 test_size=0.33 也会得到相同结果,因为 12 的 33% 大约是 4。
这两个例子还有一个非常重要的区别:你现在每次运行函数都会得到相同的结果。这是因为你通过 random_state=4 固定了随机数生成器。
下图展示了调用 train_test_split() 时发生的情况:

数据集的样本被随机打乱,然后根据你定义的大小划分为训练集和测试集。
你可以看到 y 有六个 0 和六个 1。然而,测试集中四个项目中有三个是 0。如果你希望(近似)在训练集和测试集中保持 y 值的比例,那么可以传入 stratify=y。这将启用分层拆分:
>>> x_train, x_test, y_train, y_test = train_test_split(
... x, y, test_size=0.33, random_state=4, stratify=y
... )
>>> x_train
array([[21, 22],
[ 1, 2],
[15, 16],
[13, 14],
[17, 18],
[19, 20],
[23, 24],
[ 3, 4]])
>>> x_test
array([[11, 12],
[ 7, 8],
[ 5, 6],
[ 9, 10]])
>>> y_train
array([1, 0, 1, 0, 1, 0, 0, 1])
>>> y_test
array([0, 0, 1, 1])
现在 y_train 和 y_test 中 0 和 1 的比例与原始 y 数组相同。
在某些情况下(例如对不平衡数据集进行分类时),分层拆分是可取的。不平衡数据集是指属于不同类别的样本数量存在显著差异的数据集。
最后,你可以通过 shuffle=False 关闭数据打乱和随机拆分:
>>> x_train, x_test, y_train, y_test = train_test_split(
... x, y, test_size=0.33, shuffle=False
... )
>>> x_train
array([[ 1, 2],
[ 3, 4],
[ 5, 6],
[ 7, 8],
[ 9, 10],
[11, 12],
[13, 14],
[15, 16]])
>>> x_test
array([[17, 18],
[19, 20],
[21, 22],
[23, 24]])
>>> y_train
array([0, 1, 1, 0, 1, 0, 0, 1])
>>> y_test
array([1, 0, 1, 0])
现在你得到了一个拆分结果:原始 x 和 y 数组中的前三分之二样本被分配给训练集,最后三分之一被分配给测试集。没有打乱,没有随机性。
在监督式机器学习中使用 train_test_split()
现在是时候看看 train_test_split() 在解决监督学习问题中的实际应用了。你将从一个可以用线性回归解决的小型回归问题开始,然后再看一个更大的问题。你还会看到 train_test_split() 也可用于分类。
线性回归的极简示例
在这个例子中,你将应用迄今为止所学的知识来解决一个小型回归问题。你将学习如何创建数据集、将其拆分为训练和测试子集,并用它们进行线性回归。
像往常一样,你首先导入必要的包、函数或类。你需要 NumPy、LinearRegression 和 train_test_split():
>>> import numpy as np
>>> from sklearn.linear_model import LinearRegression
>>> from sklearn.model_selection import train_test_split
现在你已经导入了所有需要的内容,可以创建两个小数组 x 和 y 来表示观测值,然后像之前一样将它们拆分为训练集和测试集:
>>> x = np.arange(20).reshape(-1, 1)
>>> y = np.array([5, 12, 11, 19, 30, 29, 23, 40, 51, 54, 74,
... 62, 68, 73, 89, 84, 89, 101, 99, 106])
>>> x
array([[ 0],
[ 1],
[ 2],
[ 3],
[ 4],
[ 5],
[ 6],
[ 7],
[ 8],
[ 9],
[10],
[11],
[12],
[13],
[14],
[15],
[16],
[17],
[18],
[19]])
>>> y
array([ 5, 12, 11, 19, 30, 29, 23, 40, 51, 54, 74, 62, 68,
73, 89, 84, 89, 101, 99, 106])
>>> x_train, x_test, y_train, y_test = train_test_split(
... x, y, test_size=8, random_state=0
... )
你的数据集包含 20 个观测值(即 x-y 对)。你指定了参数 test_size=8,因此数据集被划分为包含 12 个观测值的训练集和包含 8 个观测值的测试集。
现在你可以使用训练集来拟合模型:
>>> model = LinearRegression().fit(x_train, y_train)
>>> model.intercept_
np.float64(3.1617195496417523)
>>> model.coef_
array([5.53121801])
LinearRegression 创建表示模型的对象,而 .fit() 训练(或拟合)模型并返回它。在线性回归中,拟合模型意味着确定回归线的最佳截距(model.intercept_)和斜率(model.coef_)值。
虽然你可以使用 x_train 和 y_train 检查拟合优度,但这并不是最佳实践。模型预测性能的无偏估计应基于测试数据:
>>> model.score(x_train, y_train)
0.9868175024574795
>>> model.score(x_test, y_test)
0.9465896927715023
.score() 返回传递数据的决定系数(R²)。其最大值为 1。R² 值越高,拟合效果越好。在此例中,训练数据产生的系数略高。然而,使用测试数据计算的 R² 是模型预测性能的无偏度量。
图形显示如下:

绿色圆点表示用于训练的 x-y 对。黑色线条(称为估计回归线)由模型拟合结果(截距和斜率)定义,因此仅反映绿色圆点的位置。
白色圆点表示测试集。你使用它们来估计模型(回归线)在未用于训练的数据上的性能。
回归示例
现在你已准备好拆分更大的数据集来解决回归问题。你将使用 sklearn 中包含的 加州住房数据集(California Housing dataset)。该数据集包含 20640 个样本、8 个输入变量和房屋价值作为输出。你可以通过 sklearn.datasets.fetch_california_housing() 获取它。
首先导入 train_test_split() 和 fetch_california_housing():
>>> from sklearn.datasets import fetch_california_housing
>>> from sklearn.model_selection import train_test_split
现在你已经导入了这两个函数,可以获取你要处理的数据:
>>> x, y = fetch_california_housing(return_X_y=True)
如你所见,带参数 return_X_y=True 的 fetch_california_housing() 返回一个包含两个 NumPy 数组的元组:
- 一个包含输入的二维数组
- 一个包含输出的一维数组
下一步是像之前一样拆分数据:
>>> x_train, x_test, y_train, y_test = train_test_split(
... x, y, test_size=0.4, random_state=0
... )
现在你有了训练集和测试集。训练数据包含在 x_train 和 y_train 中,而测试数据包含在 x_test 和 y_test 中。
处理大型数据集时,通常更方便以比例形式传递训练集或测试集的大小。test_size=0.4 表示大约 40% 的样本将被分配给测试数据,其余 60% 分配给训练数据。
最后,你可以使用训练集(x_train 和 y_train)拟合模型,并使用测试集(x_test 和 y_test)对模型进行无偏评估。在此示例中,你将应用三种著名的回归算法来创建拟合数据的模型:
- 使用
LinearRegression()的线性回归 - 使用
GradientBoostingRegressor()的梯度提升 - 使用
RandomForestRegressor()的随机森林
过程与之前的示例基本相同:
- 导入所需的类。
- 使用这些类创建模型实例。
- 使用
.fit()和训练集拟合模型实例。 - 使用
.score()和测试集评估模型。
以下是针对所有三种回归算法的代码:
>>> from sklearn.linear_model import LinearRegression
>>> model = LinearRegression().fit(x_train, y_train)
>>> model.score(x_train, y_train)
0.6105322745695656
>>> model.score(x_test, y_test)
0.5982535501446862
>>> from sklearn.ensemble import GradientBoostingRegressor
>>> model = GradientBoostingRegressor(random_state=0).fit(x_train, y_train)
>>> model.score(x_train, y_train)
0.8083859166342285
>>> model.score(x_test, y_test)
0.7802104901623703
>>> from sklearn.ensemble import RandomForestRegressor
>>> model = RandomForestRegressor(random_state=0).fit(x_train, y_train)
>>> model.score(x_train, y_train)
0.9727449572570027
>>> model.score(x_test, y_test)
0.7933138227558006
你已使用训练集和测试集拟合了三个模型并评估了它们的性能。使用 .score() 获得的准确度度量是决定系数。它可以用训练集或测试集计算。但如前所述,使用测试集获得的分数代表了性能的无偏估计。
正如文档中提到的,你可以为 LinearRegression()、GradientBoostingRegressor() 和 RandomForestRegressor() 提供可选参数。GradientBoostingRegressor() 和 RandomForestRegressor() 使用 random_state 参数的原因与 train_test_split() 相同:处理算法中的随机性并确保可复现性。
对于某些方法,你可能还需要特征缩放。在这种情况下,你应该使用训练数据拟合缩放器,并用它们转换测试数据。
分类示例
你可以像处理回归分析一样使用 train_test_split() 来解决分类问题。在机器学习中,分类问题涉及训练模型为输入值打标签(或分类),并将数据集分类到不同类别中。
在《Python 中的逻辑回归》教程中,你会找到一个手写识别任务的示例。该示例提供了另一种拆分数据为训练集和测试集以避免评估过程中出现偏差的演示。
其他验证功能
sklearn.model_selection 包提供了许多与模型选择和验证相关的功能,包括:
- 交叉验证(Cross-validation)
- 学习曲线(Learning curves)
- 超参数调优(Hyperparameter tuning)
交叉验证是一组技术,通过组合预测性能度量来获得更准确的模型估计。
一种广泛使用的交叉验证方法是 k 折交叉验证(k-fold cross-validation)。在该方法中,你将数据集划分为 k 个(通常为 5 或 10)大小相等的子集(或折),然后执行 k 次训练和测试过程。每次使用不同的折作为测试集,其余所有折作为训练集。这提供了 k 个预测性能度量,然后你可以分析它们的均值和标准差。
你可以使用 KFold、StratifiedKFold、LeaveOneOut 以及 sklearn.model_selection 中的其他几个类和函数来实现交叉验证。
学习曲线(有时称为训练曲线)显示了训练集和验证集的预测分数如何随训练样本数量变化。你可以使用 learning_curve() 获取这种依赖关系,这有助于你找到训练集的最佳大小、选择超参数、比较模型等。
超参数调优(也称为超参数优化)是确定最佳超参数集以定义机器学习模型的过程。sklearn.model_selection 为此提供了多种选项,包括 GridSearchCV、RandomizedSearchCV、validation_curve() 等。数据拆分在超参数调优中也很重要。
结论
你现在知道了为什么以及如何使用 sklearn 中的 train_test_split()。你已经了解到,为了对机器学习模型的预测性能进行无偏估计,你应该使用未用于模型拟合的数据。这就是为什么你需要将数据集拆分为训练集、测试集,有时还需要验证集。
在本教程中,你学会了如何:
- 使用
train_test_split()获取训练集和测试集 - 使用
train_size和test_size参数控制子集大小 - 使用
random_state参数确定拆分的随机性 - 使用
stratify参数获得分层拆分 - 将
train_test_split()作为监督式机器学习流程的一部分使用
你还了解到 sklearn.model_selection 模块提供了多种其他模型验证工具,包括交叉验证、学习曲线和超参数调优。