使用tfds ImageFolder加载数据，VGG16调用fit时输入形状不兼容报错

错误根源 #

你的Keras模型期望输入包含批量维度（形状为(None, 363, 360, 3)，其中None代表可变的批量大小），但当前训练数据集输出的是单张图像的形状(363, 360, 3)，缺少批量维度，导致输入与模型层不匹配。

修复方案 #

在加载数据集后，添加batch()方法指定批量大小，同时搭配prefetch()优化训练效率：

train_ds, test_ds = builder.as_dataset(split=['train','test'], shuffle_files=True, as_supervised=True,)

# 设置批量大小并添加预取优化
BATCH_SIZE = 32  # 根据显存容量调整合适值
train_ds = train_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
test_ds = test_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

关键说明 #

• batch()会将多个样本打包成一个批量，自动为输入添加批量维度，完全匹配模型的输入要求。
• prefetch(tf.data.AUTOTUNE)让TensorFlow在训练当前批量的同时，提前准备下一个批量的数据，有效提升训练速度。

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