使用 scikit-learn 的 train_test_split() 分割数据集

28 Feb 2025 | 10分钟阅读

sklearn.model_selection.train_test_split() 函数用于将数据分成训练集和测试集。但是,在使用它之前,数据必须先分为输入数据集 (X) 和输出数据集 (y)。该函数返回四个子集:本研究中使用的数据集是 X_train 的训练数据集、X_test 的测试数据集、y_train 的训练标签集和 y_test 的测试标签集。将清晰矩阵划分为 X_train 和 y_train 子集以训练和拟合模型,但使用 X_test 和 y_test 来测试模型的预测。您可以设置训练数据和测试数据的数量,也可以让程序自动设置,或者定义更多的训练数据会更好。

训练集:这是提供给模型的数据,以便模型根据这些数据构建模型。

测试集:这种类型的数据集与训练集完全不同,用于模型的整体性能评估。

验证集:训练数据的子集,用于参数微调和模型准确性评估。

欠拟合:当模型由于未能理解分别定义输入和输出的质量和特性而在训练数据和新生成数据上表现不佳时,就会发生这种情况。

过拟合:当模型对训练数据集过度优化,从而在其他不同数据集上测试时表现不佳时,这很常见。

语法

交叉验证函数单独 sklearn.cross_validation.train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)

参数

*arrays:接受列表、NumPy 数组、scipy 稀疏矩阵或 pandas DataFrames。

test_size:可以是 int 或 float。如果它是 float,则表示分配给测试的数据比例(介于 0.0 和 1.0 之间)。如果它是 int,则指定测试样本的数量。如果未提供 train_size,则默认值为 0.25。

train_size:一个 int 或 float,用于确定用于训练的数据比例。

random_state:这控制数据的洗牌以获得可重现的结果。应提供一个整数以获得一致的输出。

shuffle:一个布尔值,默认为 True,表示是否在拆分之前洗牌数据。如果设置为 False,则 stratify 必须为 None。

stratify:如果指定,这将确保类分布在拆分中保持一致。默认值为 None。

返回值

一个包含用于训练和测试的数据集的列表。

拆分数据集的步骤

步骤 1:导入所需的包

在此步骤中,我们将必要的包和模块导入到 Python 环境中。

步骤 2:加载数据集

我们使用 pd.read_csv() 加载数据集,然后使用 .shape 属性检查 DataFrame 的形状。

步骤 3:分配特征和目标变量

在此阶段,我们将自变量分配给 X,将因变量分配给 y。

步骤 4:将数据拆分为训练集和测试集

我们使用 sklearn.model_selection 中的 train_test_split() 函数将数据集划分为训练集和测试集。test_size 参数指定分配给测试集的数据部分,而 random_state 确保我们的结果可以重现。

示例

输出

(20,3)
(15, 2) (5, 2)
(15,) (5,)

示例概述

然后,我们加载“carprices.csv”,为了确认数据,我们得到文件的形状是 (20, 3)。属性“Mileage”和“Age(yrs)”投影到向量 X 上,预测变量——“Sell Price($)”投影到向量 y 上。接下来,为了获得训练集和测试集,我们使用 sklearn 中的 train_test_split() 函数。因此,当 test_size 为 0.25 时,25% 的数据集(即 5 行)分配给测试集,而其余 75%(即 15 行)用于训练。random_state 意味着每次运行代码时都可以获得相同的结果。在训练之前设置的数据集称为训练数据集,用于评估模型性能的数据集称为测试数据集。

示例

输出

(20, 3)
Head of the dataframe:
   Mileage  Age(yrs)  Sell Price($)
0    69000         6          18000
1    35000         3          34000
2    57000         5          26100
3    22500         2          40000
4    46000         4          31500
Index(['Mileage', 'Age(yrs)', 'Sell Price($)'], dtype='object')
X_train:
    Mileage  Age(yrs)
13    58780         4
9     67000         6
11    79000         7
10    83000         7
7     72000         6
(15, 2)

X_test:
    Mileage  Age(yrs)
19    52000         5
18    87600         8
16    28000         2
12    59000         5
2     57000         5
(5, 2)

y_train:
13    27500
9     22000
11    19500
10    18700
7     19300
Name: Sell Price($), dtype: int64
(15,)

y_test:
19    28200
18    12800
16    35500
12    26000
2     26100
Name: Sell Price($), dtype: int64
(5,)

说明

在此示例中,执行以下步骤

导入包

此处还导入了 Numpy、pandas 和 train_test_split(作为 sklearn 的拆分),以加载、操作和拆分数据。

加载数据集

carprices.csv 数据集使用 pd.read_csv() 函数读取到 pandas 对象中,然后显示该 carprices.csv 数据集的维度。head() 最初应用参数设置为水平,以显示数据的前五行,同时显示列以确保数据类型正确构建。

创建特征和目标变量

我们选择的自变量是 Mileage 和 Age(yrs),并将其分配给 X,而我们的目标变量 Sell Price($) 分配给 y。这些将稍后用于训练和预测。

拆分数据集

train_test_split() 用于将数据拆分为训练集和测试集。在这种情况下,test_size = 0.25,因此 25% 的数据用于测试,75% 用于训练。random_state = 104 保证每次脚本运行时测试集和训练集都相同。

打印训练集和测试集

代码首先显示 X_train、X_test、y_train 和 y_test 数据集的前几条记录以及它们的维度,以验证拆分。

它有助于通过使数据可用于模型训练、模型评估或使用机器学习模型执行预测等用途来清理数据。

数据划分在模型评估中的作用

许多有监督机器学习都致力于创建能够有效地将函数输入渲染为函数输出的模型。

然而,模型评估是由问题的性质决定的。在回归任务中,确定系数或 R2 指标、均方根误差 (RMSE) 和平均绝对误差 (MAE) 是评估模型性能的一些方法。在分类问题中,估计准确率、精确度、召回率和 F1 分数是主要衡量标准。

这些指标的可接受阈值因领域而异,可从 Statistics By Jim 和 Quora 等来源获得更多见解。

而且,重要的是要注意,除非对详细指标进行清晰无偏的评估,并且这与您试图衡量模型预测准确性以及模型本身的性能所面临的挑战相关,否则这些都无法有效应用。

使用用于训练模型的数据来评估模型是不正确的,因为这会产生偏差结果。相反,应使用潜在数据来测试模型的性能。这可以通过首先将数据集划分为训练样本和测试样本来适当使用数据集来实现。

为了说明 train_test_split() 函数的目的,它将采用两个数组,例如 x 和 y,并将它们分成四个数组,如下所示

x_train - 这是用于训练的第一个数组 (x) 的部分。

x_test - 这是用于测试的第一个数组 (x) 的部分。

y_train - 这是用于训练的第二个数组 (y) 的部分。

y_test - 这是用于测试的第二个数组 (y) 的部分。

由于该函数默认以随机顺序拆分数据,因此每次运行该函数时输出不保证相同。这可能并非总是理想的情况,特别是当用户特别希望为了可重现性而获得相同输出时。

为了确保这种情况在实践中不会发生,用户可以将 random_state 参数设置为任何非负整数。random_state 参数由于其简单性而适用于大多数目的,尽管 numpy.random.RandomState 可用于更复杂的问题。

训练集、验证集和测试集

将数据集划分为子集以减少模型预测能力评估中的任何偏差非常重要。通常,数据集会随机划分为三个部分

  1. 训练集:它是将用于训练或拟合模型的数据部分。训练集对于找到最佳权重或系数至关重要,具体取决于线性或逻辑回归甚至神经网络。
  2. 验证集:有助于在调整超参数时公平地评估模型。它用于尝试打开或关闭不同的特征,例如,更改神经网络层中的单元数或支持向量机使用哪个核函数。模型在上述训练集上针对每个配置进行训练,并在验证集上评估训练后的模型。
  3. 测试集:它提供了最终模型过拟合程度的独立值。但它绝不应用于训练或验证。

在更简单的上下文中,在没有任何超参数优化要求的情况下,数据集可以仅划分为两部分:训练集和测试集。

欠拟合和过拟合

数据集拆分还有助于识别另外两个方面:欠拟合和过拟合。

  1. 欠拟合:模型甚至没有学习数据中的基本模式,它可能只是一个拟合到非线性数据集的线性回归模型。在这些情况下,劣质模型在训练数据和保留的测试数据上都表现出低性能。
  2. 过拟合:当模型变得过于复杂,以至于它对数据以及数据中存在的噪声都非常理想时,就会发生这种情况。此类模型在用于训练它们的数据上会显示出非常高的准确性,但对于从相同分布中提取的新数据将完全无用。

使用 train_test_split() 的先决条件

为了有效地评估您的模型并检测欠拟合或过拟合,您需要正确拆分数据集。现在,让我们逐步了解如何使用 scikit-learn 执行此拆分。

本教程中使用的算法是 scikit-learn network 1.5.0。Scikit-learn 是一个丰富的发行版,包含许多主要用于数据科学和机器学习的全球有用资源。本文特别关注 model_selection 模块和更高级的功能,即 train_test_split() 函数。

请注意:本教程以 1.5.0 版本作为参考材料,但所研究的功能是库的核心,通常在其他版本中以相同的方式运行。

如果您使用的是 Anaconda,您可能已经安装了 sklearn。如果您想要全新的设置,可以通过以下方式进行安装

注意:NumPy 无需额外安装,因为它与 sklearn 打包在一起。对于那些想重新学习 NumPy 的人,请查看 NumPy 教程:Python 中数据科学的第一步。

使用 train_test_split() 函数

要使用 train_test_split() 函数,必须首先将其与 Numpy 一起导入。这可以在 Jupyter notebook 或 Python 的 REPL 会话中完成

现在,让我们研究这个函数如何帮助我们将数据拆分为训练集和测试集,只需一次函数调用。它可以在数组、数据框或类似结构上执行,前提是它们的大小相等。

arrays:用户输入列表、NumPy 数组或包含数据集的 DataFrames(例如输入 x 和输出 y)。

选项

可以修改的一些关键字参数

  1. train_size:表示训练集大小。如果它是 float,则它将介于 0.0 到 1.0 之间;如果它是整数,则它将是训练样本大小。默认值为 None。
  2. test_size:表示测试集的大小。默认情况下,如果 train_size 和 test_size 都不存在,则后者将占用输入的 25%。
  3. random_state:通过随机化控制实现结果的可重现性。它允许整数或 RandomState 实例。
  4. shuffle:一个布尔标志,指示是否在拆分之前洗牌数据(默认值为 True)。
  5. stratify:如果给定数组类对象,它将导致分层拆分。

结论

为了进行拆分,本分析使用 scikit-learn 中的 train_test_split() 将数据集拆分为训练集和测试集。它在评估模型性能时很有用,因为它允许数据分区,通常可以最大限度地减少偏差,并防止过拟合。它使用户能够设置拆分比率、洗牌数据,甚至可以使用 random_state 来再次创建相同的操作。这些步骤包括加载数据集、定义特征和标签以及拆分数据和显示拆分集。

除了指定为训练集、测试集和交叉验证集之外,数据集还可以细分为包含超参数验证。当模型在训练集上表现良好但在其他集上表现不佳时,通常称为过拟合,而如果模型甚至没有捕获重要特征,则称为欠拟合。