如何将Keras MNIST数据集划分为(0 1)等5个指定分组？

1. 加载并预处理MNIST数据集 #

先通过Keras内置接口加载数据集，同时做基础的像素归一化处理：

import numpy as np
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten

# 加载原始数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()
# 将像素值归一化到0-1区间，提升模型训练稳定性
x_train = x_train / 255.0
x_test = x_test / 255.0

2. 按目标分组筛选数据并训练模型 #

定义你需要的5个数字分组，遍历每个分组筛选对应样本，再训练独立的模型：

# 设定目标数字分组
digit_groups = [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9)]

# 逐个处理每个分组
for group in digit_groups:
    # 筛选训练集中属于当前分组的样本
    train_mask = np.logical_or(y_train == group[0], y_train == group[1])
    x_group_train = x_train[train_mask]
    y_group_train = y_train[train_mask]
    
    # 筛选测试集中属于当前分组的样本
    test_mask = np.logical_or(y_test == group[0], y_test == group[1])
    x_group_test = x_test[test_mask]
    y_group_test = y_test[test_mask]
    
    # 将原标签转换为二分类标签（把分组第一个数字设为0，第二个设为1）
    y_group_train = np.where(y_group_train == group[0], 0, 1)
    y_group_test = np.where(y_group_test == group[0], 0, 1)
    
    # 构建并训练二分类模型（此处用简单全连接网络做示例）
    model = Sequential([
        Flatten(input_shape=(28, 28)),  # 将28x28的图像展平为一维向量
        Dense(128, activation='relu'),
        Dense(1, activation='sigmoid')  # 二分类输出层
    ])
    
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',  # 二分类损失函数
                  metrics=['accuracy'])
    
    print(f"开始训练分组 {group} 的模型...")
    model.fit(x_group_train, y_group_train, epochs=5, validation_data=(x_group_test, y_group_test))
    
    # 训练完成后可按需保存模型
    # model.save(f'mnist_group_{group[0]}_{group[1]}.h5')

核心细节说明 #

• 用np.logical_or快速生成符合分组条件的样本掩码，高效筛选目标数据
• 转换标签为二分类格式，适配二分类模型的损失函数和输出层设置
• 每个分组的模型完全独立，训练、保存、推理流程互不影响

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