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.
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.
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.
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.
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.
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.
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
- How to fit list of numpy array into LSTM Neural Network?
- How to fix low volatile GPU-Util with Tensorflow-GPU and Keras?
- How to fix 'Object arrays cannot be loaded when allow_pickleFalse' in the sketch_rnn algorithm
- How to force tensorflow tensors to be symmetric?
- How to Fine tune existing Tensorflow Object Detection model to recognize additional classes?
- How to Fine tune existing Tensorflow Object Detection model to recognize additional classes?
- How to force tensorflow to use all available GPUs?
- How to freeze weights in certain layer with Keras?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.