基于Python与TensorFlow实现无磁盘副本的数据集扩增方法

需求说明 #

需要训练一个基于ResNet50的0-4类图像分类模型，现有约2万张自制图像数据集。要求不生成磁盘物理副本，将数据集扩增2倍、4倍甚至10倍，让模型在训练时能看到原图及多个不同的扩增变体（每个原图对应若干个扩增版本）。

原数据集结构：

0: 0.jpg, 1.jpg, 2.jpg, ....
1: 0.jpg, 1.jpg, 2.jpg, ....
2: 0.jpg, 1.jpg, 2.jpg, ....
3: 0.jpg, 1.jpg, 2.jpg, ....
...

期望模型看到的训练样本结构：

0: 0.jpg, 0_copy0.jpg, 0_copy1.jpg, ....
1: 0.jpg, 0_copy0.jpg, 0_copy1.jpg, ....
2: 0.jpg, 0_copy0.jpg, 0_copy1.jpg, ....
3: 0.jpg, 0_copy0.jpg, 0_copy1.jpg, ....
...

原代码问题分析 #

原代码存在两处关键问题：

1. 无效的数据集修改循环：for idx, img in enumerate(train_ds): train_ds[idx] = data_augmentation 无法修改tf.data.Dataset对象，属于无效操作，需删除。
2. 未实现固定倍数扩增：当前仅在每个batch上随机应用扩增，但无法保证每个原图被模型看到多次（对应不同变体），无法达到“扩增N倍”的效果。

修改后的完整代码 #

import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import load_model

# 配置参数
DATASET_PATH = "MY_DATASET"
CLASS_NAMES = ["0", "1", "2", "3", "4"]
img_height, img_width = 180, 180
batch_size = 32
# 设置扩增倍数，支持2、4、10等自定义数值
AUGMENTATION_TIMES = 4
AUTOTUNE = tf.data.AUTOTUNE

# 定义数据扩增流水线
data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip("horizontal"),
    tf.keras.layers.RandomRotation(0.2),
    tf.keras.layers.RandomContrast(0.3),
    tf.keras.layers.RandomBrightness(0.3)
])

# 加载原始数据集（训练集+验证集）
train_ds_raw = tf.keras.preprocessing.image_dataset_from_directory(
    DATASET_PATH,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    DATASET_PATH,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

# 定义数据集扩增函数：生成指定倍数的变体并拼接
def augment_dataset(ds, augment_times):
    # 保留原始图像
    augmented_ds = ds
    # 循环生成augment_times-1次扩增后的数据集并拼接
    for _ in range(augment_times - 1):
        augmented_ds = augmented_ds.concatenate(
            ds.map(lambda x, y: (data_augmentation(x, training=True), y),
                   num_parallel_calls=AUTOTUNE)
        )
    # 打乱顺序+预取，提升训练效率
    return augmented_ds.shuffle(buffer_size=1000).prefetch(buffer_size=AUTOTUNE)

# 对训练集执行指定倍数的扩增
train_ds = augment_dataset(train_ds_raw, AUGMENTATION_TIMES)

# 构建ResNet50模型
resnet_model = Sequential()
pretrained_model = tf.keras.applications.ResNet50(
    include_top=False,
    input_shape=(180, 180, 3),
    pooling='avg',
    classes=len(CLASS_NAMES),
    weights='imagenet'
)
# 冻结预训练层，避免破坏已学习的特征
for layer in pretrained_model.layers:
    layer.trainable = False

resnet_model.add(pretrained_model)
resnet_model.add(Flatten())
resnet_model.add(Dense(512, activation='relu'))
resnet_model.add(Dense(5, activation='softmax'))

# 编译模型
resnet_model.compile(
    optimizer=Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 训练模型
epochs = 10  # 可根据实际需求调整轮次
history = resnet_model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=epochs
)

# 保存训练好的模型
resnet_model.save("LOCAL_PATH/output.h5", save_format="h5")

# 绘制精度曲线
plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.axis(ymin=0, ymax=1)
plt.grid()
plt.title('模型精度')
plt.ylabel('精度')
plt.xlabel('训练轮次')
plt.legend(['训练集', '验证集'])
plt.show()

# 绘制损失曲线
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.grid()
plt.title('模型损失')
plt.ylabel('损失')
plt.xlabel('训练轮次')
plt.legend(['训练集', '验证集'])
plt.show()

关键修改说明 #

1. 扩增倍数可控：通过AUGMENTATION_TIMES参数直接设置扩增倍数，比如设为4时，每个原图会生成1张原图+3张不同的扩增变体，总数据集规模变为原来的4倍。
2. 无物理副本生成：所有扩增操作在内存中实时完成，通过数据集拼接实现“虚拟扩增”，完全不需要在磁盘上生成图像副本。
3. 训练效率优化：添加shuffle打乱数据集顺序，prefetch实现数据预取，避免训练过程中出现等待数据的情况。
4. 移除无效代码：删除了原代码中无法修改数据集的循环操作，精简了冗余逻辑。

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