使用train_test_split划分数据集

目录

1.简介

2.参数说明

3.返回值

4.示例

---

1.简介

  train_test_split 是 sklearn 库中的一个函数，用于将数据集分割为训练集和测试集。

2.参数说明

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

• *arrays：一个或多个数组样例，可以是NumPy数组、Pandas DataFrame或任何其他序列类型。这些数组将根据指定的比例被分割。

• test_size：单个浮点数，表示测试集占总数据的比例。如果设置，train_size 将被忽略。默认值为0.25，即测试集占25%。

• train_size：单个浮点数或整数，表示训练集占总数据的比例或具体的样本数量。如果设置了 test_size，则此参数将被忽略。

• random_state：随机数生成器的种子。如果为整数，random_state 将固定随机数生成器的状态，这有助于确保结果的可重复性。如果为 None（默认值），则每次运行结果可能不同。

• shuffle：布尔值，是否在分割前对数据进行随机打乱。默认为 True，意味着数据将在分割前被打乱。如果数据已经预先排序或有特定的顺序，可以设置为 False。

• stratify：当数据集是分层的时（例如，在分类问题中，每个类别的分布是已知的），stratify 参数可以用来确保分割后的训练集和测试集中各类别的比例与原始数据集中的比例相同。这通常用于保持类别的分布一致性，提高模型评估的有效性。

3.返回值

  train_test_split 返回两个或四个数组，具体取决于输入的数组数量。对于两个数组，返回的是训练集和测试集的分割。如果输入了更多的数组，将按照相同的分割比例分割所有数组。

4.示例

from sklearn.model_selection import train_test_split
import numpy as np

#  X 是特征数据，y 是标签数据
X = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16], [17, 18], [19, 20]])
y = np.array([0, 1, 0, 1, 1, 0, 1, 1, 0, 0])

# 将数据集分割为训练集和测试集，测试集占30%
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)

print(X_train)
print(X_test)
print(y_train)
print(y_test)

输出：