sklearn train_test_split - Python

在机器学习中，我们通常需要使用训练集和测试集来构建和评估模型。sklearn提供了train_test_split函数，可以帮助我们随机拆分数据集，生成训练集和测试集。

安装

要使用train_test_split，需要先安装sklearn包。在终端中输入以下命令进行安装：

pip install -U scikit-learn

用法

train_test_split的使用非常简单。以下是一些常用参数：

train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)

其中，*arrays表示输入的数据集，可以是一个或多个数组；test_size表示测试集的比例；train_size表示训练集的比例（如果和test_size都未提供，则默认80%训练集和20%测试集）；random_state表示随机种子（如果指定了随机种子，每次生成的数据集都是固定的）；shuffle表示是否要打乱数据集；stratify表示按照什么因素进行分层抽样（通常在类别不平衡的情况下使用）。

以下是一个使用train_test_split的示例：

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

# 加载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("训练集大小：", X_train.shape, y_train.shape)
print("测试集大小：", X_test.shape, y_test.shape)

输出结果如下：

训练集大小： (105, 4) (105,)
测试集大小： (45, 4) (45,)

注意事项

在使用train_test_split时，需要注意以下几点：

• 训练集和测试集的比例通常根据数据集大小和模型复杂度来调整；
• 随机种子可以保证每次生成的数据集是固定的，方便进行比较和调试；
• 分层抽样可以解决类别不平衡问题，但在数据量较小时可能会导致分层效果不佳；
• 由于train_test_split是随机拆分数据集的，因此拆分结果可能会受到随机性影响，需要进行多次拆分并取均值来获得更稳定的结果。