Train multi-class image classifier in Keras

Keras

Image Classification

Multi-class Classification

Deep Learning

Machine Learning

Train multi-class image classifier in Keras

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Deep Learning TensorFlow & Keras Computer Vision Machine Learning

Introduction

Training a multi-class image classifier in Keras is mostly about keeping the data pipeline, label encoding, and model output consistent. The network architecture matters, but many training failures come from mismatched class labels, missing normalization, or weak validation setup rather than from the CNN itself. A solid baseline should be simple, reproducible, and easy to improve later with transfer learning.

Organize the Dataset Correctly

Keras works well with a directory-per-class layout:

text

1data/
2  train/
3    class_a/
4    class_b/
5    class_c/
6  val/
7    class_a/
8    class_b/
9    class_c/

This lets you use the built-in dataset loader:

python

1import tensorflow as tf
2
3IMG_SIZE = (224, 224)
4BATCH_SIZE = 32
5
6train_ds = tf.keras.utils.image_dataset_from_directory(
7    "data/train",
8    image_size=IMG_SIZE,
9    batch_size=BATCH_SIZE,
10    seed=42,
11)
12
13val_ds = tf.keras.utils.image_dataset_from_directory(
14    "data/val",
15    image_size=IMG_SIZE,
16    batch_size=BATCH_SIZE,
17    seed=42,
18    shuffle=False,
19)
20
21class_names = train_ds.class_names
22print(class_names)

By default, labels are integer class IDs. That choice affects the correct loss and metric configuration later.

Normalize and Prepare the Input Pipeline

Images should be normalized before training. A simple Rescaling layer is enough for a baseline model.

python

1AUTOTUNE = tf.data.AUTOTUNE
2
3normalizer = tf.keras.layers.Rescaling(1.0 / 255.0)
4
5train_ds = train_ds.map(lambda x, y: (normalizer(x), y), num_parallel_calls=AUTOTUNE)
6val_ds = val_ds.map(lambda x, y: (normalizer(x), y), num_parallel_calls=AUTOTUNE)
7
8train_ds = train_ds.prefetch(AUTOTUNE)
9val_ds = val_ds.prefetch(AUTOTUNE)

If the dataset fits in memory, cache() can speed things up further. If not, skip it and keep the pipeline streaming.

Build a Simple Baseline CNN

Start with a small model you can reason about before jumping to more complex transfer-learning setups.

python

1num_classes = len(class_names)
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(224, 224, 3)),
5    tf.keras.layers.Conv2D(32, 3, activation="relu"),
6    tf.keras.layers.MaxPooling2D(),
7    tf.keras.layers.Conv2D(64, 3, activation="relu"),
8    tf.keras.layers.MaxPooling2D(),
9    tf.keras.layers.Conv2D(128, 3, activation="relu"),
10    tf.keras.layers.GlobalAveragePooling2D(),
11    tf.keras.layers.Dropout(0.3),
12    tf.keras.layers.Dense(num_classes, activation="softmax"),
13])
14
15model.compile(
16    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
17    loss="sparse_categorical_crossentropy",
18    metrics=["accuracy"],
19)
20
21model.summary()

The final dense layer size must equal the number of classes, and softmax pairs naturally with multi-class classification.

Add Data Augmentation Carefully

Augmentation can improve generalization, but keep it realistic for the domain.

python

1data_augmentation = tf.keras.Sequential([
2    tf.keras.layers.RandomFlip("horizontal"),
3    tf.keras.layers.RandomRotation(0.05),
4    tf.keras.layers.RandomZoom(0.1),
5])

You can insert it near the model input:

python

1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(224, 224, 3)),
3    data_augmentation,
4    tf.keras.layers.Rescaling(1.0 / 255.0),
5    tf.keras.layers.Conv2D(32, 3, activation="relu"),
6    tf.keras.layers.MaxPooling2D(),
7    tf.keras.layers.Conv2D(64, 3, activation="relu"),
8    tf.keras.layers.GlobalAveragePooling2D(),
9    tf.keras.layers.Dense(num_classes, activation="softmax"),
10])

For some tasks, such as medical imaging or OCR, careless augmentation can damage label meaning. Use domain judgment.

Train with Callbacks

Callbacks help stop overtraining and preserve the best model.

python

1callbacks = [
2    tf.keras.callbacks.EarlyStopping(
3        monitor="val_loss",
4        patience=5,
5        restore_best_weights=True
6    ),
7    tf.keras.callbacks.ModelCheckpoint(
8        "best_model.keras",
9        monitor="val_accuracy",
10        save_best_only=True
11    ),
12]
13
14history = model.fit(
15    train_ds,
16    validation_data=val_ds,
17    epochs=30,
18    callbacks=callbacks,
19)

That gives a reliable training loop without much extra complexity.

Evaluate Beyond Accuracy

Validation accuracy is helpful, but you should also inspect actual predictions and class-specific errors.

python

1val_loss, val_acc = model.evaluate(val_ds)
2print("val_loss:", val_loss)
3print("val_acc:", val_acc)
4
5for images, labels in val_ds.take(1):
6    probs = model.predict(images, verbose=0)
7    preds = tf.argmax(probs, axis=1)
8    print("true:", labels[:10].numpy())
9    print("pred:", preds[:10].numpy())

For imbalanced datasets, add confusion matrices or per-class precision and recall. A model can show acceptable overall accuracy while failing badly on one class.

Upgrade Path: Transfer Learning

If the baseline underperforms, use a pretrained backbone such as EfficientNet or MobileNet. That usually improves results on small or medium datasets.

python

1base = tf.keras.applications.MobileNetV2(
2    include_top=False,
3    weights="imagenet",
4    input_shape=(224, 224, 3),
5    pooling="avg"
6)
7base.trainable = False

Then place your classifier head on top. This is often a better starting point than trying to deepen a small custom CNN endlessly.

Common Pitfalls

The most common mistake is mismatching labels, loss, and output layer. Integer labels with categorical_crossentropy, or one-hot labels with sparse_categorical_crossentropy, will produce wrong training behavior.

Another issue is forgetting normalization or using inconsistent preprocessing between training and inference.

Developers also often rely on training accuracy alone. If validation metrics are ignored, the model can look good while overfitting badly.

Summary

Keep class directories, label encoding, and final output size consistent.
Normalize images and build a reproducible tf.data pipeline.
Start with a simple baseline CNN before moving to transfer learning.
Use callbacks to stop early and preserve the best weights.
Inspect validation behavior and per-class errors, not just headline accuracy.

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Train multi-class image classifier in Keras

ML System Design practice on Codemia

Introduction

Organize the Dataset Correctly

Normalize and Prepare the Input Pipeline

Build a Simple Baseline CNN

Add Data Augmentation Carefully

Train with Callbacks

Evaluate Beyond Accuracy

Upgrade Path: Transfer Learning

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia