如何在TensorFlow中创建自定义数据集并替换MNIST训练神经网络？

我正在尝试让神经网络按需求训练。目前用mnist.npz数据集训练手写数字识别模型，但想切换为自定义数据集却无从下手，多种尝试均失败。我的训练准备代码如下：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.datasets import mnist
from tensorflow.keras import backend as K
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

def show_min_max(array, i):
  random_image = array[i]
  print(random_image.min(), random_image.max())

def plot_image(array, i, labels):
  plt.imshow(np.squeeze(array[i]))
  plt.title(" Class " + str(labels[i]))
  plt.xticks([])
  plt.yticks([])
  plt.show()

img_rows, img_cols = 28, 28  
num_classes = 10 

(train_images, train_labels), (test_images, test_labels) = mnist.load_data(path="mnist.npz")  #training data
(train_images_backup, train_labels_backup), (test_images_backup, test_labels_backup) = mnist.load_data(path="mnist.npz") #untrained backup data

print(train_images.shape) 
print(test_images.shape) 

train_images = train_images.reshape(train_images.shape[0],  img_rows, img_cols, 1) 
test_images = test_images.reshape(test_images.shape[0], img_rows, img_cols, 1) 
input_shape = (img_rows, img_cols, 1) 

plot_image(train_images, 1, train_labels)
show_min_max(train_images, 100) 

train_images[3000:]=255-train_images[3000:] 

train_images = train_images.astype('float32') 
test_images = test_images.astype('float32') 

train_images /= 255 
test_images /= 255 

plot_image(train_images, 100, train_labels) 
show_min_max(train_images, 100) 

train_labels = keras.utils.to_categorical(train_labels, num_classes) 
test_labels = keras.utils.to_categorical(test_labels, num_classes) 

已尝试方法：

• 查阅TensorFlow与MNIST代码，想指向自定义数据集但需将文件夹转为npz格式；
• 使用image_dataset_from_directory()函数未达预期；
• 将图片转为npz文件仍无效。

我不清楚mnist.load_data()加载的x_train、y_train、x_test、y_test含义，也不知如何将文件夹数据集转为npz文件，希望得到帮助，可指导转格式或从基础代码重新梳理。

一、先搞懂mnist.load_data()的返回值 #

mnist.load_data()返回包含两个元组的元组：

• 第一个元组(train_images, train_labels)：
  • train_images：形状为(60000, 28, 28)的numpy数组，对应6万张28x28的灰度训练图片，像素值范围0-255；
  • train_labels：形状为(60000,)的numpy数组，对应每张训练图片的数字标签（0-9）。
• 第二个元组(test_images, test_labels)：
  • test_images：形状为(10000, 28, 28)的numpy数组，对应1万张测试图片；
  • test_labels：形状为(10000,)的numpy数组，对应测试图片的标签。

简单来说，x_train/x_test是图片像素数据，y_train/y_test是对应的分类标签。

二、方法1：将自定义图片数据集转为npz格式 #

假设你的自定义数据集文件夹结构为（按数字分类存储）：

custom_dataset/
    train/
        0/
            img0.png
            img1.png
            ...
        1/
            ...
        ...
        9/
    test/
        0/
            ...
        ...
        9/

步骤1：读取图片并转为numpy数组 #

import os
import numpy as np
from PIL import Image

def load_custom_dataset(data_dir, img_size=(28,28)):
    images = []
    labels = []
    # 遍历每个数字分类文件夹
    for label in sorted(os.listdir(data_dir)):
        label_dir = os.path.join(data_dir, label)
        if not os.path.isdir(label_dir):
            continue
        # 读取文件夹内所有图片
        for img_name in os.listdir(label_dir):
            img_path = os.path.join(label_dir, img_name)
            try:
                # 转为灰度图并调整尺寸至28x28
                img = Image.open(img_path).convert('L').resize(img_size)
                img_array = np.array(img)
                images.append(img_array)
                labels.append(int(label))
            except Exception as e:
                print(f"跳过损坏图片: {img_path}, 错误: {e}")
    # 转为numpy数组返回
    return np.array(images), np.array(labels)

# 加载训练和测试集
train_images, train_labels = load_custom_dataset("custom_dataset/train")
test_images, test_labels = load_custom_dataset("custom_dataset/test")

# 验证形状是否符合预期
print(f"训练集形状: {train_images.shape}, 标签形状: {train_labels.shape}")
print(f"测试集形状: {test_images.shape}, 标签形状: {test_labels.shape}")

步骤2：保存为npz文件 #

# 保存时使用与MNIST一致的键名，方便直接替换原代码
np.savez("custom_mnist.npz", x_train=train_images, y_train=train_labels, x_test=test_images, y_test=test_labels)

步骤3：替换原代码加载逻辑 #

直接替换原代码中mnist.load_data()的行即可，后续预处理代码无需修改：

(train_images, train_labels), (test_images, test_labels) = mnist.load_data(path="custom_mnist.npz")

三、方法2：直接用image_dataset_from_directory()加载（无需转npz） #

之前使用该函数未成功，大概率是预处理逻辑有误，以下是正确用法：

步骤1：加载并预处理数据集 #

import tensorflow as tf
from tensorflow import keras

img_size = (28, 28)
batch_size = 32
num_classes = 10

# 加载训练集
train_ds = tf.keras.utils.image_dataset_from_directory(
    "custom_dataset/train",
    image_size=img_size,
    color_mode="grayscale",  # 转为灰度图，匹配MNIST格式
    batch_size=batch_size,
    label_mode="int"
)

# 加载测试集
test_ds = tf.keras.utils.image_dataset_from_directory(
    "custom_dataset/test",
    image_size=img_size,
    color_mode="grayscale",
    batch_size=batch_size,
    label_mode="int"
)

# 数据标准化（等价于原代码中的/255）
normalization_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))
test_ds = test_ds.map(lambda x, y: (normalization_layer(x), y))

# 可选：实现原代码中的图片反转逻辑（示例为随机反转50%图片）
def invert_images(images, labels):
    inverted = 1.0 - images
    mask = tf.random.uniform(shape=tf.shape(labels), minval=0, maxval=2, dtype=tf.int32) == 1
    images = tf.where(mask[:, tf.newaxis, tf.newaxis, tf.newaxis], inverted, images)
    return images, labels

train_ds = train_ds.map(invert_images)

# 转独热编码（匹配原代码的标签格式）
train_ds = train_ds.map(lambda x, y: (x, tf.one_hot(y, num_classes)))
test_ds = test_ds.map(lambda x, y: (x, tf.one_hot(y, num_classes)))

步骤2：用数据集训练模型 #

训练时直接传入数据集即可，无需修改原模型结构：

model.fit(train_ds, epochs=10, validation_data=test_ds)

这种方法无需将所有图片加载到内存，适合大规模数据集，且流程更简洁。

四、验证数据集加载正确性 #

不管用哪种方法，都可以用原代码中的工具函数验证：

• 方法1（npz）：直接调用原代码的plot_image和show_min_max函数；
• 方法2（dataset）：可先转numpy数组再验证：

# 从数据集中取出一个批次
for images, labels in train_ds.take(1):
    img = images[0].numpy()
    # 从独热编码转回数字标签
    label = np.argmax(labels[0].numpy())
    # 可视化
    plt.imshow(np.squeeze(img))
    plt.title(f"Class {label}")
    plt.xticks([])
    plt.yticks([])
    plt.show()
    print(f"像素值范围: {img.min()} - {img.max()}")

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