使用tf.data加载MNIST时CNN模型精度低、训练慢的问题排查

我是机器学习和Python新手，抱歉提出入门级问题。

我正在使用MNIST数据集学习简单CNN模型，首先将MNIST数据保存为PNG文件：

import cv2
import os

from keras.datasets import mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

img_root = './images'
train = os.path.join(img_root, 'train')
test = os.path.join(img_root, 'test')

if not os.path.exists(img_root):
    os.mkdir(img_root)

if not os.path.exists(train):
    os.mkdir(train)

if not os.path.exists(test):
    os.mkdir(test)

# Save Train images
for i in range(x_train.shape[0]):
    img_dir = os.path.join(train, str(y_train[i]))
    if not os.path.exists(img_dir):
        os.mkdir(img_dir)
    image_out = os.path.join(img_dir, str(i) + ".png")
    cv2.imwrite(image_out, x_train[i])

# Save Test images
for i in range(x_test.shape[0]):
    img_dir = os.path.join(test, str(y_test[i]))
    if not os.path.exists(img_dir):
        os.mkdir(img_dir)
    image_out = os.path.join(img_dir, str(i) + ".png")
    cv2.imwrite(image_out, x_test[i])

随后我采用三种不同方式加载这些图像并训练测试模型：

方法1 “cv_only”：使用cv2.imread将所有图像加载到numpy数组，不使用tf.data #

结果如下：

Epoch 48/50
469/469 [==============================] - 2s 5ms/step - loss: 0.0025 - accuracy: 0.9992 - val_loss: 0.0490 - val_accuracy: 0.9937

Epoch 49/50
469/469 [==============================] - 2s 5ms/step - loss: 0.0042 - accuracy: 0.9990 - val_loss: 0.0477 - val_accuracy: 0.9924

Epoch 50/50
469/469 [==============================] - 2s 5ms/step - loss: 0.0033 - accuracy: 0.9991 - val_loss: 0.0485 - val_accuracy: 0.9914

Learning time: 199[s]

方法2 “tf_cv”：使用tf.data.Dataset.from_tensor_slices()、map()和batch()构建数据集，map方法内使用cv2.imread读取图像 #

结果如下：

Epoch 48/50
469/469 [==============================] - 32s 68ms/step - loss: 5.8638e-07 - accuracy: 1.0000 - val_loss: 4.7863 - val_accuracy: 0.4507

Epoch 49/50
469/469 [==============================] - 32s 68ms/step - loss: 4.6416e-07 - accuracy: 1.0000 - val_loss: 4.8436 - val_accuracy: 0.4514

Epoch 50/50
469/469 [==============================] - 32s 69ms/step - loss: 3.6748e-07 - accuracy: 1.0000 - val_loss: 4.8742 - val_accuracy: 0.4517

Learning time: 1638[s]

方法3 “tf_only”：与方法2相同使用tf.data构建数据集，但map方法内使用tf.io.read_file读取图像 #

结果如下：

Epoch 48/50
469/469 [==============================] - 16s 33ms/step - loss: 8.5148e-07 - accuracy: 1.0000 - val_loss: 5.4608 - val_accuracy: 0.4065

Epoch 49/50
469/469 [==============================] - 16s 34ms/step - loss: 6.7230e-07 - accuracy: 1.0000 - val_loss: 5.4721 - val_accuracy: 0.4085

Epoch 50/50
469/469 [==============================] - 15s 33ms/step - loss: 5.3065e-07 - accuracy: 1.0000 - val_loss: 5.4845 - val_accuracy: 0.4087

Learning time: 700[s]

请问我在方法2和3中哪里出错了？恳请帮助。

以下是我使用的完整模型代码：

import numpy as np
import cv2
import glob
import datetime

import tensorflow as tf
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, BatchNormalization
from keras.layers import Dense, Flatten
from keras.utils import to_categorical

# Valid Options are "cv_only", "tf_cv" and "tf_only"
data_load_method = "tf_cv"

train_images = "./images/train/"
test_images = "./images/test/"

# network parameters
batch_size = 128

categories = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
num_classes = len(categories)


def make_sample_imgs(files):
    global X, Y
    X = []
    Y = []
    for cat, fname in files:
        X.append(cv2.imread(fname, cv2.IMREAD_GRAYSCALE))
        Y.append(cat)
    return np.array(X), np.array(Y)


def make_sample(files):
    global X, Y
    X = []
    Y = []
    for cat, fname in files:
        X.append(fname)
        Y.append(cat)
    return np.array(X), np.array(Y)


def process_path_cv2(image_path, label):

    def load_image(path):
        image_gray = cv2.imread(path.decode("utf-8"), cv2.IMREAD_GRAYSCALE)
        image_gray = image_gray.astype('float32')/255
        image_gray = tf.expand_dims(image_gray, axis=-1)
        return image_gray

    image = tf.numpy_function(load_image, [image_path], tf.float32)
    return image, label


def process_path_tf(image_path, label):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_image(image, channels=1)
    image = tf.image.convert_image_dtype(image, tf.float32)
    return image, label


start_time = datetime.datetime.now()

train = []
test = []

for idx, cat in enumerate(categories):
    image_dir = train_images + cat
    files = glob.glob(image_dir + "/*.png")
    for f in files:
        train.append((idx, f))
    print("Train ", image_dir, "append ", len(files), "files!")

for idx, cat in enumerate(categories):
    image_dir = test_images + cat
    files = glob.glob(image_dir + "/*.png")
    for f in files:
        test.append((idx, f))
    print("Test ", image_dir, "append ", len(files), "files!")


if data_load_method == "cv_only":
    x_train, y_train = make_sample_imgs(train)
    x_test, y_test = make_sample_imgs(test)

    # convert to one-hot vector
    y_train = to_categorical(y_train)
    y_test = to_categorical(y_test)

    x_train = x_train.reshape(-1, 28, 28, 1)
    x_test = x_test.reshape(-1, 28, 28, 1)

    x_train = x_train.astype('float32') / 255
    x_test = x_test.astype('float32') / 255

elif data_load_method == "tf_cv" or data_load_method == "tf_only":
    x_train, y_train = make_sample(train)
    x_test, y_test = make_sample(test)

    # convert to one-hot vector
    y_train = to_categorical(y_train)
    y_test = to_categorical(y_test)

    train_datasets = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    test_datasets = tf.data.Dataset.from_tensor_slices((x_test, y_test))

    if data_load_method == "tf_cv":
        train_datasets = train_datasets.map(process_path_cv2, num_parallel_calls=tf.data.AUTOTUNE)
        test_datasets = test_datasets.map(process_path_cv2, num_parallel_calls=tf.data.AUTOTUNE)

    else:
        train_datasets = train_datasets.map(process_path_tf, num_parallel_calls=tf.data.AUTOTUNE)
        test_datasets = test_datasets.map(process_path_tf, num_parallel_calls=tf.data.AUTOTUNE)

    # train_datasets = train_datasets.cache()
    # test_datasets = test_datasets.cache()

    train_datasets = train_datasets.batch(batch_size)
    test_datasets = test_datasets.batch(batch_size)
    train_datasets = train_datasets.prefetch(tf.data.AUTOTUNE)
    test_datasets = test_datasets.prefetch(tf.data.AUTOTUNE)

else:
    print("Method not defined!")
    exit()


model = Sequential()

model.add(Conv2D(filters=64, kernel_size=(3, 3), padding="same", activation="relu", input_shape=(28, 28, 1)))
model.add(Conv2D(filters=64, kernel_size=(3, 3), padding="same", strides=2, activation="relu"))
model.add(BatchNormalization())

model.add(Conv2D(filters=128, kernel_size=(3, 3), padding="same", strides=2, activation="relu"))
model.add(Conv2D(filters=128, kernel_size=(3, 3), padding="same", strides=2, activation="relu"))
model.add(BatchNormalization())

model.add(Conv2D(filters=256, kernel_size=(3, 3), padding="same", strides=2, activation="relu"))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(BatchNormalization())

model.add(Flatten())
model.add(Dense(512, activation="relu"))

model.add(Dense(10, activation="softmax"))

model.summary()

model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])

if data_load_method == "cv_only":
    model.fit(x_train, y_train, epochs=50, batch_size=batch_size, validation_data=(x_test, y_test))

elif data_load_method == "tf_cv" or data_load_method == "tf_only":
    model.fit(train_datasets, epochs=50, batch_size=batch_size, validation_data=test_datasets)

end_time = datetime.datetime.now()
time_diff = (end_time - start_time)
learning_time = time_diff.total_seconds()
print(f'Learning time: {int(learning_time)}[s]')

如果我使用：

train_datasets = train_datasets.cache()
test_datasets = test_datasets.cache()

训练时间可降至约150秒，但精度仍会下降。我使用的环境为Python 3.10、Keras 2.10.0、TensorFlow-GPU 2.10.0。

抱歉代码有些混乱，我仍在学习编程。

核心问题1：训练数据未打乱导致严重过拟合 #

方法1中model.fit()默认会在每个epoch打乱训练数据，而你用tf.data构建的数据集没有添加shuffle操作。你的训练数据是按类别顺序存储的（先存所有0，再存所有1，以此类推），模型会快速记住这种顺序，导致训练精度瞬间拉满，但完全无法泛化到测试集，出现训练loss趋近于0、验证loss极高的情况。

修复：给训练数据集添加shuffle #

在构建train_datasets时，加上shuffle()，建议设置一个足够大的buffer_size（比如训练集总数60000）：

train_datasets = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_datasets = train_datasets.shuffle(60000)  # 添加这一行

核心问题2：图像预处理的一致性与Tensor形状推断问题 #

方法3的tf图像读取问题 #

tf.image.decode_image不会固定返回的张量形状，可能导致模型接收的输入形状不一致。需要强制将图像形状固定为(28,28,1)，修改process_path_tf：

def process_path_tf(image_path, label):
    image = tf.io.read_file(image_path)
    # 改用decode_png，明确指定形状
    image = tf.image.decode_png(image, channels=1)
    image = tf.reshape(image, (28, 28, 1))  # 强制固定形状
    image = tf.image.convert_image_dtype(image, tf.float32)
    return image, label

方法2的cv2读取问题 #

用tf.numpy_function包装cv2读取时，TensorFlow无法自动推断返回张量的形状，需要手动设置形状：

def process_path_cv2(image_path, label):
    def load_image(path):
        image_gray = cv2.imread(path.decode("utf-8"), cv2.IMREAD_GRAYSCALE)
        image_gray = image_gray.astype('float32')/255
        return image_gray.reshape(28,28,1)  # 先转成正确形状再返回

    image = tf.numpy_function(load_image, [image_path], tf.float32)
    image.set_shape((28,28,1))  # 手动设置张量形状
    return image, label

次要问题：batch_size重复设置导致的潜在问题 #

你已经用train_datasets.batch(batch_size)设置了批次大小，在model.fit()中不需要再传入batch_size=batch_size，否则可能引发冲突（TensorFlow会忽略fit中的batch_size，但最好去掉避免混淆）：

# 原代码
model.fit(train_datasets, epochs=50, batch_size=batch_size, validation_data=test_datasets)
# 修改后
model.fit(train_datasets, epochs=50, validation_data=test_datasets)

效率优化：保留cache操作 #

你已经发现cache()可以大幅提升训练速度，建议保留这个操作，它会把预处理后的图像缓存到内存/磁盘，避免每个epoch重复读取和处理图像：

train_datasets = train_datasets.map(...)
train_datasets = train_datasets.cache()  # 放在map之后，batch之前
train_datasets = train_datasets.batch(batch_size)
train_datasets = train_datasets.prefetch(tf.data.AUTOTUNE)

修复后的预期效果 #

• 训练精度不会瞬间到100%，会逐步上升，最终和方法1接近
• 验证精度会提升到和方法1相当的水平（99%左右）
• 训练时间在启用cache后会和方法1接近甚至更快

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