Fine-tuning strategy in PyTorch - Model Pipeline Trace

Model Pipeline - Fine-tuning strategy

This pipeline shows how a pre-trained model is adapted to a new task by fine-tuning. We start with a model trained on a large dataset, then adjust its weights slightly using new task data to improve performance.

Data Flow - 6 Stages

1Load pre-trained model

N/A→Load a model trained on a large dataset (e.g., ImageNet)→Model with fixed architecture and pre-trained weights

ResNet50 model with weights trained on ImageNet

↓

2Replace final layer

Model with output layer for 1000 classes→Replace last layer to match new task classes (e.g., 10 classes)→Model with output layer for 10 classes

Replace ResNet50's 1000-class layer with a 10-class layer

↓

3Freeze base layers

Model with all layers trainable→Freeze weights of all layers except the new final layer→Model with only final layer trainable

Freeze convolutional layers, train only classifier layer

↓

4Train final layer

Training data: 5000 images x 3 channels x 224 x 224→Train only the final layer on new data→Updated final layer weights

Train classifier on 5000 labeled images for 5 epochs

↓

5Unfreeze some layers

Model with final layer trained→Unfreeze last few layers to fine-tune→Model with last layers trainable

Unfreeze last 2 convolutional blocks

↓

6Fine-tune entire model

Training data: 5000 images x 3 channels x 224 x 224→Train unfrozen layers with a low learning rate→Fine-tuned model weights

Train for 10 epochs with learning rate 0.0001

Training Trace - Epoch by Epoch

Loss
1.2 |*       
1.0 | *      
0.8 |  *     
0.6 |   *    
0.4 |    *   
0.2 |       *
    +--------
     1 3 5 10
     Epochs

Epoch	Loss ↓	Accuracy ↑	Observation
1	1.2	0.55	Training only final layer starts with moderate loss and accuracy
3	0.8	0.70	Loss decreases and accuracy improves as final layer learns
5	0.6	0.78	Final layer training converges with good accuracy
6	0.58	0.80	Unfreeze last layers and continue training with low learning rate
8	0.45	0.85	Fine-tuning improves model performance further
10	0.40	0.88	Loss decreases steadily, accuracy reaches high level

Prediction Trace - 5 Layers

Layer 1: Input image

Layer 2: Feature extraction layers (frozen)

Layer 3: New classifier layer (trainable)

Layer 4: Softmax activation

Layer 5: Prediction

Model Quiz - 3 Questions

Why do we freeze most layers at the start of fine-tuning?

A Because frozen layers improve accuracy automatically

B To speed up training by removing layers

C To keep learned features and train only the new task-specific layer

D To prevent the model from making any predictions

Key Insight

Fine-tuning leverages existing knowledge in a pre-trained model by first training only new layers, then carefully adjusting deeper layers with a low learning rate. This approach helps the model adapt to new tasks efficiently while preserving useful features.