如何使用TensorFlow正确将图像数据集按8:1:1比例拆分

原代码问题说明 #

• 第一种image_dataset_from_directory拆分逻辑错误：两次调用分别拆分了总数据的80%为训练集、10%为验证集，剩余10%总数据没有被定义为测试集，也没有做显式拆分。
• 第二种ImageDataGenerator实现的错误：
  1. 三个迭代器均未设置validation_split和subset参数，实际加载的都是全量数据集，完全没有拆分，会导致严重的数据泄露。
  2. class_mode设置为binary（对应二分类任务），但你的分类数为26类，且损失函数为多分类用的SparseCategoricalCrossentropy，参数完全不匹配。
• 模型形状异常的问题：没有在模型第一层指定输入形状，导致summary输出维度全为None。

正确实现代码 #

方案1：使用image_dataset_from_directory（推荐，执行效率更高） #

import tensorflow as tf

# 自定义参数配置
img_height = 32
img_width = 32
batch_size = 32
dataset_dir = "你的数据集目录路径"

# 第一步：拆分总数据的80%为训练集
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
  dataset_dir,
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size,
  color_mode="grayscale" # 对应灰度图需求，RGB图可删除此行
)

# 第二步：取出剩余20%的总数据作为验证+测试的合集
val_test_ds = tf.keras.preprocessing.image_dataset_from_directory(
  dataset_dir,
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size,
  color_mode="grayscale"
)

# 第三步：将20%的合集按1:1拆分，得到各占总数据10%的验证集、测试集
val_test_batches = tf.data.experimental.cardinality(val_test_ds).numpy()
val_ds = val_test_ds.take(val_test_batches // 2)
test_ds = val_test_ds.skip(val_test_batches // 2)

方案2：修复后的ImageDataGenerator实现 #

from tensorflow.keras.preprocessing.image import ImageDataGenerator

img_size = (32, 32)
batch_size = 32
dataset_dir = "你的数据集目录路径"

# 先设置全局拆分比例
datagen = ImageDataGenerator(validation_split=0.2) # 先拆分20%为验证+测试合集

# 训练集：占总数据80%
train_it = datagen.flow_from_directory(
  dataset_dir, 
  target_size=img_size, 
  color_mode='grayscale', 
  class_mode='sparse', # 多分类配合SparseCategoricalCrossentropy用sparse模式
  batch_size=batch_size, 
  shuffle=True, 
  subset="training"
)

# 验证+测试合集：占总数据20%
val_test_it = datagen.flow_from_directory(
  dataset_dir, 
  target_size=img_size, 
  color_mode='grayscale', 
  class_mode='sparse',
  batch_size=batch_size, 
  shuffle=True, 
  subset="validation"
)

# 手动拆分合集为1:1的验证、测试集
val_test_samples = val_test_it.samples
val_samples = val_test_samples // 2
test_samples = val_test_samples - val_samples

val_it = val_test_it
val_it.samples = val_samples
val_it.n = val_samples

test_it = val_test_it
test_it.samples = test_samples
test_it.n = test_samples
test_it.index_array = val_test_it.index_array[val_samples:]

模型修复 #

在模型的第一个预处理层指定输入形状，即可解决summary输出维度全为None的问题：

num_classes = 26
img_height = 32
img_width = 32

model = tf.keras.Sequential([
  # 灰度图通道数为1，RGB图改为3即可
  tf.keras.layers.experimental.preprocessing.Rescaling(1./255, input_shape=(img_height, img_width, 1)),
  tf.keras.layers.Conv2D(32, 3, activation='relu'),
  tf.keras.layers.MaxPooling2D(),
  tf.keras.layers.Conv2D(32, 3, activation='relu'),
  tf.keras.layers.MaxPooling2D(),
  tf.keras.layers.Conv2D(32, 3, activation='relu'),
  tf.keras.layers.MaxPooling2D(),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dense(num_classes)
])
model.compile(
  optimizer='adam',
  loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
  metrics=['accuracy']
)

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