scikit-learn库学习之train_test_split函数
scikit-learn库学习之train_test_split函数
一、简介
train_test_split是scikit-learn库中用于将数据集拆分为训练集和测试集的函数。它在进行模型训练和评估时非常有用，通过将数据集划分为独立的训练和测试部分，可以有效评估模型的性能和泛化能力。
二、语法和参数
sklearn.model_selection.train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)
*arrays: 待划分的数组或矩阵。test_size: float 或 int, 可选，测试集占总样本的比例（0到1之间）或样本数量。默认值为0.25。train_size: float 或 int, 可选，训练集占总样本的比例（0到1之间）或样本数量。若为None，则为补齐test_size。random_state: int 或 RandomState instance, 可选，随机数生成器的种子。shuffle: boolean, 可选，是否在划分之前对数据进行洗牌。默认值为True。stratify: array-like, 可选，按照该参数进行分层采样。
三、实例
3.1 基本用法
import numpy as np
from sklearn.model_selection import train_test_split
# 创建示例数据
X, y = np.arange(10).reshape((5, 2)), np.array([0, 1, 0, 1, 0])
# 拆分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 输出结果
print("X_train:\n", X_train)
print("X_test:\n", X_test)
print("y_train:\n", y_train)
print("y_test:\n", y_test)
输出：
X_train:
[[8 9]
[4 5]
[0 1]
[6 7]]
X_test:
[[2 3]]
y_train:
[0 0 0 1]
y_test:
[1]
3.2 按比例划分数据集
import numpy as np
from sklearn.model_selection import train_test_split
# 创建示例数据
X, y = np.arange(20).reshape((10, 2)), np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
# 按0.3的比例划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
# 输出结果
print("X_train:\n", X_train)
print("X_test:\n", X_test)
print("y_train:\n", y_train)
print("y_test:\n", y_test)
输出：
X_train:
[[ 0 1]
[14 15]
[ 4 5]
[18 19]
[ 8 9]
[ 6 7]
[12 13]]
X_test:
[[16 17]
[ 2 3]
[10 11]]
y_train:
[0 1 0 1 0 1 0]
y_test:
[0 1 1]
3.3 按类别比例划分数据集
import numpy as np
from sklearn.model_selection import train_test_split
# 创建示例数据
X, y = np.arange(20).reshape((10, 2)), np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
# 按类别比例划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
# 输出结果
print("X_train:\n", X_train)
print("X_test:\n", X_test)
print("y_train:\n", y_train)
print("y_test:\n", y_test)
输出：
X_train:
[[ 6 7]
[16 17]
[ 4 5]
[14 15]
[10 11]
[ 2 3]
[ 8 9]]
X_test:
[[ 0 1]
[12 13]
[18 19]]
y_train:
[1 0 0 1 1 1 0]
y_test:
[0 0 1]
四、注意事项
- 确保
shuffle参数为True以确保数据的随机性，除非有特殊需要不进行洗牌。 stratify参数用于确保划分后每个子集中的类别比例与原始数据集一致。- 设置
random_state参数以保证每次划分数据的一致性。 - 如果指定了
test_size，train_size会自动补齐为1减去test_size的值。
原文地址：https://blog.csdn.net/qq_46396470/article/details/140537374
免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！