机器学习：训练集与测试集分割train_test_split

原创已于 2024-01-22 19:17:56 修改 · 1.3w 阅读

90 ·

CC 4.0 BY-SA版权

文章标签：

#机器学习 #人工智能 #深度学习 #sklearn

于 2023-08-01 16:17:30 首次发布

深度学习之pytorch 专栏收录该内容

23 篇文章

订阅专栏

本文详细介绍了scikit-learn库中的train_test_split函数，用于Python机器学习中数据集的划分，包括参数说明、使用示例以及sklearn库在数据预处理和机器学习任务中的重要性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1 引言

在使用机器学习训练模型算法的过程中，为提高模型的泛化能力、防止过拟合等目的，需要将整体数据划分为训练集和测试集两部分，训练集用于模型训练，测试集用于模型的验证。此时，使用train_test_split函数可便捷高效的实现数据训练集与测试集的划分。

2 train_test_split介绍

train_test_split函数来自scikit-learn库（也称为sklearn），安装命令：

pip install sklearn

函数的导入：

from sklearn.model_selection import train_test_split

2.1 函数定义

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

2.2 参数说明

*arrays: 单个数组或元组，表示需要划分的数据集。如果传入多个数组，则必须保证每个数组的第一维大小相同。
test_size: 测试集的大小（占总数据集的比例，值为0.0-1.0，表示测试集占总样本比例）。默认值为0.25，即将传入数据的25%作为测试集。
train_size: 训练集的大小（占总数据集的比例，值为0.0-1.0，表示训练集占总样本比例）。默认值为None，此时和test_size互补，即训练集的大小为(1-test_size)。
random_state: 随机数种子。可以设置一个整数，用于复现结果。默认为None。其实是该组随机数的编号，在需要重复试验的时候，保证得到一组一样的随机数。（比如每次都填1，其他参数一样的情况下你得到的随机数组是一样的。但填0或不填，每次都会不一样。）
shuffle: 是否随机打乱数据。默认为True。
stratify: 可选参数，用于进行分层抽样。传入标签数组，保证划分后的训练集和测试集中各类别样本比例与原始数据集相同。默认为None，即普通的随机划分。（此参数作用是保持测试集与整个数据集里的数据分类比例一致，比如有1000个数据，800个属于A类，200个属于B类。设置stratify = y_lable，test_size=0.25，split之后数据组成如下：training: 750个数据，其中600个属于A类，150个属于B类；testing: 250个数据，其中200个属于A类，50个属于B类）

2.3 返回值

该函数返回一个元组(X_train, X_test, y_train, y_test)，其中X_train表示训练集的特征数据，X_test表示测试集的特征数据，y_train表示训练集的标签数据，y_test表示测试集的标签数据。

2.4 注意事项

test_size和train_size必须至少有一个设置为非None。
当传入多个数组时，请确保每个数组的第一维大小相同。
random_state要设置一个整数值，从而保证每次获取相同的训练集和测试集
当使用分层抽样时，请确保传入的标签数组是正确的。

3 train_test_split使用

3.1 使用train_test_split分割Iris数据

from sklearn import datasets
from sklearn.model_selection import train_test_split

# 加载Iris数据集
iris = datasets.load_iris()
X = iris.data
y = iris.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=1)
train_dataset = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
test_dataset = TensorDataset(torch.from_numpy(X_test), torch.from_numpy(y_test ))

batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

print(X_train)
print(X_test)

结果展示：

X_train=[[6.5 2.8 4.6 1.5]
 [6.7 2.5 5.8 1.8]
 [6.8 3.  5.5 2.1]
 [5.1 3.5 1.4 0.3]
 [6.  2.2 5.  1.5]
 ......此处数据省略
 [4.9 3.6 1.4 0.1]]
X_test=[[5.8 4.  1.2 0.2]
 [5.1 2.5 3.  1.1]
 [6.6 3.  4.4 1.4]
 [5.4 3.9 1.3 0.4]
 [7.9 3.8 6.4 2. ]
 ......此处数据省略
 [5.2 3.4 1.4 0.2]]

3.2 使用train_test_split分割水果识别数据

在/opt/dataset下存放着水果图片的分类数据文件夹（文件夹名称为标签），每个文件夹下存储着多张对应标签的水果图片，如下所示：

以apple文件夹为例，图片内容如下：

数据加载和分割数据集的代码如下：

from torchvision.datasets import ImageFolder
from sklearn.model_selection import train_test_split

# 图像变换
transform = transforms.Compose([transforms.Resize((224, 224)),
                                transforms.ToTensor(),
                                transforms.Normalize(
                                     mean=[0.5, 0.5, 0.5],
                                     std=[0.5, 0.5, 0.5]
                                ), ])
# 加载数据集
dataset = ImageFolder('/opt/dataset', transform=transform)

# 划分训练集与测试集
train_dataset, valid_dataset = train_test_split(dataset, test_size=0.2, random_state=10)

batch_size = 64
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
test_loader = DataLoader(dataset=valid_dataset, batch_size=batch_size, shuffle=True, drop_last=True)

4 sklearn介绍

train_test_split（）是sklearn
.cross_validation模块中用来随机划分训练集和测试集。

sklearn，全称scikit-learn，是python中的机器学习库，建立在numpy、scipy、matplotlib等数据科学包的基础之上，涵盖了机器学习中的样例数据、数据预处理、模型验证、特征选择、分类、回归、聚类、降维等几乎所有环节，功能十分强大。与深度学习库存在pytorch、TensorFlow等多种框架可选不同，sklearn是python中传统机器学习的首选库，不存在其他竞争者。

sklearn为初学者提供了一些经典数据集，通过这些数据集可快速搭建机器学习任务、对比模型性能。数据集主要围绕分类和回归两类经典任务，对于不同需求，常用数据集简介如下：

load_breast_cancer：乳腺癌数据集，特征为连续数值变量，标签为0或1的二分类任务
load_iris：经典鸢尾花数据集，特征为连续数值变量，标签为0/1/2的三分类任务，且各类样本数量均衡，均为50个
load_wine：红酒数据集，与鸢尾花数据集特点类似，也是用于连续特征的3分类任务，不同之处在于各类样本数量轻微不均衡
load_digits：小型手写数字数据集（之所以称为小型，是因为还有大型的手写数字数据集mnist），包含0-9共10种标签，各类样本均衡，与前面3个数据集最大不同在于特征也是离散数值0—16之间，例如在进行多项式朴素贝叶斯模型、ID3树模型时，可用该数据集
load_boston：波士顿房价数据集，连续特征拟合房价，适用于回归任务

值得指出，sklearn除了load系列经典数据集外，还支持自定义数据集make系列和下载数据集fetch系列（load系列为安装sklearn库时自带，而fetch则需额外下载），这为更多的学习任务场景提供了便利。

sklearn中的各模型均有规范的数据输入输出格式，一般以np.array和pd.dataframe为标准格式，所以一些字符串的离散标签是不能直接用于模型训练的；同时为了加快模型训练速度和保证训练精度，往往还需对数据进行预处理，例如在以距离作为度量进行训练时则必须考虑去量纲化的问题。为此，sklearn提供了一些常用的数据预处理功能，常用的包括：

MinMaxScaler：归一化去量纲处理，适用于数据有明显的上下限，不会存在严重的异常值，例如考试得分0-100之间的数据可首选归一化处理
StandardScaler：标准化去量纲处理，适用于可能存在极大或极小的异常值，此时用MinMaxScaler时，可能因单个异常点而将其他数值变换的过于集中，而用标准正态分布去量纲则可有效避免这一问题
Binarizer：二值化处理，适用于将连续变量离散化
OneHotEncoder：独热编码，一种经典的编码方式，适用于离散标签间不存在明确的大小相对关系时。例如对于民族特征进行编码时，若将其编码为0-55的数值，则对于以距离作为度量的模型则意味着民族之间存在"大小"和"远近"关系，而用独热编码则将每个民族转换为一个由1个"1"和55个"0"组成的向量。弊端就是当分类标签过多时，容易带来维度灾难，而特征又过于稀疏
Ordinary：数值编码，适用于某些标签编码为数值后不影响模型理解和训练时。例如，当民族为待分类标签时，则可将其简单编码为0-55之间的数字