Scikit Learn Train Test Split - Python

当我们使用 Scikit Learn 来训练机器学习模型的时候，我们通常需要将数据集划分为训练集和测试集。Scikit Learn 提供了一个函数来方便地进行数据集划分：train_test_split()。

函数的语法

以下是 train_test_split() 函数的语法：

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

以下是每个参数的含义：

• X：包含特征的数组或稀疏矩阵。
• y：包含标签的数组。
• test_size：用于测试集的样本占比。默认值为 0.25。
• random_state：随机数生成器的种子。

函数将返回四个数组：

• X_train：训练集特征向量。
• X_test：测试集特征向量。
• y_train：训练集标签。
• y_test：测试集标签。

基本用法

以下是 train_test_split() 函数的一个示例：

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris

iris = load_iris()

X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.3, random_state=42)

print("训练集特征向量：\n", X_train)
print("测试集特征向量：\n", X_test)
print("训练集标签：\n", y_train)
print("测试集标签：\n", y_test)

该示例从鸢尾花数据集中加载数据，并将数据集分成训练集和测试集（测试集的大小为 30%）。最后，该示例将输出划分后的训练集和测试集的特征向量和标签。

其他用法

分层抽样

train_test_split() 函数还支持分层抽样，可以将数据集中不同类别的样本数量分别分配给训练集和测试集。以下是分层抽样的示例：

X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.3, random_state=42, stratify=iris.target)

该示例使用 stratify 参数在训练集和测试集之间保持了鸢尾花数据集中各类别样本数量的相对比例。

多个数据集划分

我们还可以使用 train_test_split() 函数同时将多个数据集划分为训练集和测试集。以下是划分多个数据集的示例：

X_train1, X_test1, y_train1, y_test1, X_train2, X_test2, y_train2, y_test2 = train_test_split(iris.data, iris.target, wine.data, wine.target, test_size=0.3, random_state=42)

该示例将鸢尾花数据集和葡萄酒数据集同时划分为训练集和测试集。

总结

在本文中，我们学习了 Scikit Learn 的 train_test_split() 函数，该函数可以将数据集划分为训练集和测试集。我们还学习了 train_test_split() 函数的语法、基本用法和其他用法。现在，我们可以使用 train_test_split() 函数来方便地进行数据集划分了。