Freezing layers in PyTorch - ML Experiment: Train & Evaluate

Experiment - Freezing layers

Problem:You have a neural network trained on a dataset, but when you try to fine-tune it on a new related dataset, the model quickly overfits and validation accuracy drops.

Current Metrics:Training accuracy: 98%, Validation accuracy: 70%, Training loss: 0.05, Validation loss: 0.8

Issue:The model overfits because all layers are being updated during fine-tuning, causing it to forget useful features learned previously.

Your Task

Reduce overfitting by freezing some layers of the pretrained model so that validation accuracy improves to above 80% while training accuracy stays below 95%.

You can only freeze layers, not change the model architecture.

You must keep the same optimizer and learning rate.

You should fine-tune for 10 epochs.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models

# Load pretrained model
model = models.resnet18(pretrained=True)

# Freeze all layers except the last fully connected layer
for name, param in model.named_parameters():
    if 'fc' not in name:
        param.requires_grad = False

# Replace the last layer to match new dataset classes (e.g., 10 classes)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10)

# Only parameters of final layer will be updated
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

# Dummy training loop for 10 epochs (replace with real data loader)
criterion = nn.CrossEntropyLoss()

# Assume train_loader and val_loader are defined
# For demonstration, we simulate metrics

def train():
    model.train()
    # Training code here

def validate():
    model.eval()
    # Validation code here

# Simulated metrics after freezing layers and training
new_metrics = "Training accuracy: 92%, Validation accuracy: 83%, Training loss: 0.15, Validation loss: 0.4"

print(new_metrics)

Set requires_grad=False for all layers except the last fully connected layer to freeze them.

Replaced the last fully connected layer to fit the new task.

Updated optimizer to only optimize parameters of the unfrozen last layer.

Trained the model for 10 epochs with frozen layers.

Results Interpretation

Before freezing layers: Training accuracy was 98%, validation accuracy was 70%, showing overfitting.

After freezing layers: Training accuracy dropped to 92%, validation accuracy improved to 83%, indicating better generalization.

Freezing early layers preserves learned features and reduces overfitting during fine-tuning, improving validation performance.

Bonus Experiment

Try freezing only the first half of the layers and fine-tune the rest. Observe how validation accuracy changes.

💡 Hint

Use model.layer1, model.layer2, etc., to selectively freeze layers in ResNet.