📅  最后修改于: 2023-12-03 15:05:05.322000             🧑  作者: Mango
当我们使用 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)
以下是每个参数的含义:
函数将返回四个数组:
以下是 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() 函数来方便地进行数据集划分了。