TensorFlow
fine-tuning
pretrained models
deep learning
neural networks

How to Fine-tuning a Pretrained Network in Tensorflow?

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

Introduction

Fine-tuning a pretrained TensorFlow model is usually the fastest way to get strong results on a new dataset. The key is not to unfreeze everything immediately. A stable workflow trains a new prediction head first, then unfreezes part of the backbone with a much smaller learning rate.

Start with a Frozen Backbone

The first stage is ordinary transfer learning. Load a pretrained model, freeze it, and attach a task-specific head.

python
1import tensorflow as tf
2
3IMG_SIZE = (224, 224)
4NUM_CLASSES = 5
5
6base = tf.keras.applications.EfficientNetB0(
7    include_top=False,
8    weights="imagenet",
9    input_shape=IMG_SIZE + (3,),
10)
11base.trainable = False
12
13inputs = tf.keras.Input(shape=IMG_SIZE + (3,))
14x = tf.keras.applications.efficientnet.preprocess_input(inputs)
15x = base(x, training=False)
16x = tf.keras.layers.GlobalAveragePooling2D()(x)
17x = tf.keras.layers.Dropout(0.2)(x)
18outputs = tf.keras.layers.Dense(NUM_CLASSES, activation="softmax")(x)
19
20model = tf.keras.Model(inputs, outputs)
21model.compile(
22    optimizer=tf.keras.optimizers.Adam(1e-3),
23    loss="sparse_categorical_crossentropy",
24    metrics=["accuracy"],
25)

Calling the backbone with training=False helps keep normalization behavior stable while only the head is learning.

Build a Clean Input Pipeline

Fine-tuning often fails because of inconsistent preprocessing rather than because of model architecture. Keep preprocessing aligned with the chosen backbone and apply augmentation only to training data.

python
1def make_dataset(images, labels, training=False):
2    ds = tf.data.Dataset.from_tensor_slices((images, labels))
3    if training:
4        ds = ds.shuffle(2048, seed=42, reshuffle_each_iteration=True)
5        aug = tf.keras.Sequential([
6            tf.keras.layers.RandomFlip("horizontal"),
7            tf.keras.layers.RandomRotation(0.08),
8            tf.keras.layers.RandomZoom(0.1),
9        ])
10        ds = ds.map(lambda x, y: (aug(x, training=True), y),
11                    num_parallel_calls=tf.data.AUTOTUNE)
12    return ds.batch(32).prefetch(tf.data.AUTOTUNE)

This keeps the training and validation paths predictable.

Unfreeze Only Part of the Model

After the new head has learned useful task boundaries, unfreeze only the upper layers of the backbone and reduce the learning rate.

python
1history_head = model.fit(train_ds, validation_data=val_ds, epochs=5)
2
3base.trainable = True
4for layer in base.layers[:-30]:
5    layer.trainable = False
6
7model.compile(
8    optimizer=tf.keras.optimizers.Adam(1e-5),
9    loss="sparse_categorical_crossentropy",
10    metrics=["accuracy"],
11)
12
13callbacks = [
14    tf.keras.callbacks.EarlyStopping(
15        monitor="val_loss", patience=3, restore_best_weights=True
16    ),
17    tf.keras.callbacks.ReduceLROnPlateau(
18        monitor="val_loss", factor=0.2, patience=2
19    ),
20]
21
22history_ft = model.fit(train_ds, validation_data=val_ds, epochs=15, callbacks=callbacks)

The much smaller learning rate is essential. Without it, the fine-tuning stage can destroy useful pretrained features quickly.

Evaluate More Than Accuracy

After fine-tuning, inspect more than the top-line metric. Per-class failures often hide behind an apparently good overall accuracy score.

python
1import numpy as np
2from sklearn.metrics import classification_report, confusion_matrix
3
4pred = model.predict(test_ds)
5y_pred = np.argmax(pred, axis=1)
6y_true = np.concatenate([y.numpy() for _, y in test_ds], axis=0)
7
8print(confusion_matrix(y_true, y_pred))
9print(classification_report(y_true, y_pred))

This is especially important on imbalanced datasets, where a model can look good globally while failing minority classes.

Save the Final Artifact Clearly

Once the model is acceptable, save it with a versioned name and keep the preprocessing assumptions close to the artifact.

python
model.save("models/birds_effnetb0_v3.keras")

Reproducibility is part of fine-tuning quality, not a separate concern.

If the dataset is small, keep a close eye on validation curves during the unfreezing stage. Fine-tuning can overfit very quickly once pretrained layers become trainable, so early stopping and careful checkpoint review matter more than in the initial head-only phase.

Common Pitfalls

A common mistake is unfreezing the whole backbone immediately. That usually makes optimization unstable and can erase useful pretrained structure.

Another is forgetting to use the preprocessing function expected by the backbone architecture. Pretrained weights assume a particular input convention.

Developers also often keep the same learning rate for both training stages. Fine-tuning almost always needs a lower rate than head training.

Summary

  • Fine-tune in two stages: head training first, selective unfreezing second.
  • Keep preprocessing consistent with the chosen pretrained model.
  • Lower the learning rate significantly after unfreezing.
  • Evaluate class-level behavior, not just overall accuracy.
  • Save the final model artifact together with its preprocessing assumptions.

Related reading
Free course
Beginner
7 lessons
2 hours
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

All Rights Reserved.