如何使用scikit-learn库在Python中将数据集分为训练集和测试集？

Scikit-learn，通常称为sklearn，是Python中用于实现机器学习算法的库。它功能强大且健壮，因为它提供了各种工具来进行统计建模，包括分类、回归、聚类、降维等等，使用Python的强大和稳定的接口实现。它建立在Numpy、SciPy和Matplotlib库之上。

在将输入数据传递给机器学习算法之前，需要将其分为训练数据集和测试数据集。

输入数据集适合所选模型后，将在该模型上对其进行训练。当进行训练时，模型从数据中学习，还学习到了在新数据上进行泛化。测试数据集将不会在模型的训练过程中使用。

一旦调整好了所有的超参数，设置了最佳权重，就可以将测试数据集提供给机器学习算法。这是用于检查算法在新数据上泛化效果的数据集。让我们看一下如何使用scikit-learn库拆分数据。

更多Python相关文章，请阅读：Python 教程

示例

from sklearn.datasets import load_iris
my_data = load_iris()
X = my_data.data
y = my_data.target
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=2
)
print("训练数据的特征维度")
print(X_train.shape)
print("测试数据的特征维度")
print(X_test.shape)
print("训练数据目标值的维度")
print(y_train.shape)
print("测试数据目标值的维度")
print(y_test.shape)

输出

训练数据的特征维度
(120, 4)
测试数据的特征维度
(30, 4)
训练数据目标值的维度
(120,)
测试数据目标值的维度
(30,)

解释

• 导入所需的包。
• 将所需的数据集加载到环境中。
• 将特征和目标值从数据集中分离出来。
• 将训练和测试数据按比例分为80%和20%。
• 这意味着将使用20%的数据来检查模型在新数据上的泛化效果。
• 将这些拆分，以及数据维度打印到控制台上。