如何使用sklearn拆分.mat格式文件中的MNIST训练集与测试集

1. 导入依赖库 #

需要用到scipy读取mat文件、sklearn做数据集拆分、numpy处理数组：

import scipy.io as sio
from sklearn.model_selection import train_test_split
import numpy as np

2. 读取.mat文件并提取数据 #

# 替换为你本地.mat文件的实际存储路径
mat_data = sio.loadmat('mnist_data.mat')

如果需要将原有拆分的训练、测试数据合并后重新随机拆分，执行以下代码合并所有样本：

# 提取原数据中的特征和标签
old_train_X = mat_data['train_fea1']
old_train_y = mat_data['train_gnd1']
old_test_X = mat_data['test_fea1']
old_test_y = mat_data['test_gnd1']

# 合并所有特征样本
all_X = np.vstack((old_train_X, old_test_X))
# 把二维标签数组拉平为一维，适配sklearn接口要求
all_y = np.vstack((old_train_y, old_test_y)).ravel()

3. 用sklearn完成数据集拆分 #

合并全量数据后重新拆分训练/测试集 #

# 按8:2比例拆分训练集和测试集，参数可根据需求调整
X_train, X_test, y_train, y_test = train_test_split(
    all_X, all_y,
    test_size=0.2,
    random_state=42,
    stratify=all_y
)

参数说明：

• test_size：取值范围0~1，代表测试集占总样本的比例
• random_state：固定随机种子，保证每次运行的拆分结果一致，方便复现实验
• stratify=all_y：按标签类别比例分层拆分，保证训练集、测试集的类别分布和原数据集一致，避免分类任务出现类别偏移问题

仅拆分原有训练集为训练集+验证集 #

如果不需要打乱原有训练、测试集的划分，只想把原有训练集拆成新训练集和验证集，直接执行以下代码即可：

X_train_new, X_val, y_train_new, y_val = train_test_split(
    old_train_X, old_train_y.ravel(),
    test_size=0.25,
    random_state=42,
    stratify=old_train_y
)

内容的提问来源于stack exchange，提问作者BlueCurve