Hey there! Dealing with 70k+ images in Keras can definitely hit memory or speed bottlenecks, but let’s walk through practical, actionable fixes to get your model training smoothly:

1. Optimize Data Loading with tf.data.Dataset #

The ImageDataGenerator generator is convenient, but converting it to a tf.data.Dataset unlocks better performance optimizations like prefetching and caching, which reduce disk I/O delays and let training run in parallel with data loading:

import tensorflow as tf

# Convert your existing generators to tf.data.Dataset
train_ds = tf.data.Dataset.from_generator(
    lambda: train_generator,
    output_types=(tf.float32, tf.float32),
    output_shapes=([None, img_width, img_height, 3], [None, 15])
)
val_ds = tf.data.Dataset.from_generator(
    lambda: validation_generator,
    output_types=(tf.float32, tf.float32),
    output_shapes=([None, img_width, img_height, 3], [None, 15])
)

# Add caching and prefetching to speed up data pipeline
train_ds = train_ds.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

2. Gradient Accumulation for Small GPU Memory #

If your GPU can’t handle large batch sizes (which is common with big datasets), gradient accumulation lets you simulate a larger batch by accumulating gradients over multiple steps before updating weights:

accumulation_steps = 4  # Equivalent to batch_size * 4
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    step_count = 0
    for x_batch, y_batch in train_ds:
        with tf.GradientTape() as tape:
            y_pred = model(x_batch, training=True)
            loss = model.compiled_loss(y_batch, y_pred)
        
        gradients = tape.gradient(loss, model.trainable_variables)
        # Apply gradients only after accumulating steps
        if step_count % accumulation_steps == 0:
            model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        step_count += 1
    
    # Run validation after each epoch
    val_loss, val_acc = model.evaluate(val_ds)
    print(f"Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_acc:.4f}")

3. Integrate Data Augmentation into the Model (GPU-Accelerated) #

Instead of using ImageDataGenerator for augmentation (which runs on CPU), move augmentation to model layers so it runs on the GPU—this drastically speeds up preprocessing for large datasets:

from tensorflow.keras import layers

# Define augmentation as a reusable layer stack
data_augmentation = tf.keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.1),
    layers.RandomZoom(0.1),
])

# Build your model with augmentation included
inputs = layers.Input(shape=(img_width, img_height, 3))
x = data_augmentation(inputs)
x = layers.Rescaling(1./255)(x)  # Skip if your datagen already does rescaling
# Add your core model layers
x = layers.Conv2D(32, (3,3), activation='relu')(x)
x = layers.MaxPooling2D()(x)
# ... add more convolution/pooling layers as needed
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(15, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)

4. Mixed Precision Training #

Mixed precision uses 16-bit floats for most computations (while keeping critical operations in 32-bit) to cut memory usage and speed up training—most modern GPUs support this natively:

from tensorflow.keras.mixed_precision import set_global_policy

# Enable mixed precision
set_global_policy('mixed_float16')

# Compile your model as usual
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# Optional: Ensure final layer outputs are in 32-bit to avoid precision issues
# Replace your final Dense layer with:
# outputs = layers.Dense(15)(x)
# outputs = layers.Activation('softmax', dtype='float32')(outputs)

5. Use Transfer Learning with Lightweight Models #

Training a full CNN from scratch on 70k images is resource-heavy. Instead, use a pre-trained lightweight model like MobileNetV2 or EfficientNetB0, freeze its base layers, and only train a small classification head:

base_model = tf.keras.applications.MobileNetV2(
    input_shape=(img_width, img_height, 3),
    include_top=False,
    weights='imagenet'
)
base_model.trainable = False  # Freeze pre-trained layers

inputs = layers.Input(shape=(img_width, img_height, 3))
x = tf.keras.applications.mobilenet_v2.preprocess_input(inputs)
x = base_model(x, training=False)
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(15, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)

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

After initial training, you can unfreeze the top few layers of the base model for fine-tuning if you need better performance.

6. Convert Dataset to TFRecord Format #

If disk I/O is a bottleneck (e.g., using HDDs), converting your images to TFRecord (a binary format) reduces read times significantly:

# Example function to write images to TFRecord
def write_tfrecord(image_paths, labels, output_path):
    with tf.io.TFRecordWriter(output_path) as writer:
        for img_path, label in zip(image_paths, labels):
            img = tf.io.read_file(img_path)
            img = tf.image.decode_jpeg(img, channels=3)
            img = tf.image.resize(img, (img_width, img_height))
            feature = {
                'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[tf.io.serialize_tensor(img).numpy()])),
                'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))
            }
            example = tf.train.Example(features=tf.train.Features(feature=feature))
            writer.write(example.SerializeToString())

# Generate image paths and labels (you'll need to collect these from your directories)
# write_tfrecord(train_image_paths, train_labels, 'train.tfrecord')
# write_tfrecord(val_image_paths, val_labels, 'val.tfrecord')

# Load TFRecord dataset
def parse_tfrecord(example):
    feature_desc = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64),
    }
    example = tf.io.parse_single_example(example, feature_desc)
    image = tf.io.parse_tensor(example['image'], out_type=tf.float32)
    image.set_shape((img_width, img_height, 3))
    label = tf.one_hot(example['label'], depth=15)
    return image, label

train_ds = tf.data.TFRecordDataset('train.tfrecord').map(parse_tfrecord).batch(batch_size).prefetch(tf.data.AUTOTUNE)
val_ds = tf.data.TFRecordDataset('val.tfrecord').map(parse_tfrecord).batch(batch_size).prefetch(tf.data.AUTOTUNE)

Start with the simplest fixes first—like switching to tf.data.Dataset and using transfer learning—before moving to more involved steps like TFRecord conversion. These should help you handle the large dataset without hitting memory or speed issues.

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