Sunday, August 23, 2026
HomeData ScienceTransfer Learning – Fine-Tuning Pretrained Models in Python 2026

Transfer Learning – Fine-Tuning Pretrained Models in Python 2026

Table of Content

Training a deep learning model from scratch requires millions of labelled examples and days of GPU time. Transfer learning sidesteps this by starting with a model already trained on a large dataset (like ImageNet or Wikipedia) and fine-tuning it on your specific task. In practice, transfer learning lets you achieve state-of-the-art results with just a few hundred labelled examples and minutes of training. This guide shows you how.

How Transfer Learning Works

A neural network trained on ImageNet (1.2 million images, 1000 classes) has learned to detect edges, textures, shapes, and object parts in its early layers. These representations are broadly useful for any vision task. Transfer learning reuses these learned features for a new task — fine-tuning only the final classification layers (feature extraction) or the full network (full fine-tuning).

Feature Extraction vs Full Fine-Tuning

Feature extraction freezes all pretrained weights and only trains a new classification head. Use this when your dataset is small (<1000 samples) or very similar to the original training data. Full fine-tuning unfreezes all weights and trains the entire network on your data (usually with a low learning rate for early layers). Use this when your dataset is large enough (>10k samples) or quite different from the pretraining domain.

Image Classification with PyTorch (ResNet)

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from torch import nn, optim
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

# Load pretrained ResNet50
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# FEATURE EXTRACTION: freeze all pretrained layers
for param in model.parameters():
    param.requires_grad = False

# Replace the final FC layer for your number of classes
num_classes = 10
model.fc = nn.Linear(model.fc.in_features, num_classes)

# Only parameters in fc will be updated
optimizer = optim.Adam(model.fc.parameters(), lr=1e-3)

Full Fine-Tuning with Differential Learning Rates

# Unfreeze all layers
for param in model.parameters():
    param.requires_grad = True

# Use much lower LR for early layers (they already have good features)
optimizer = optim.Adam([
    {'params': model.layer1.parameters(), 'lr': 1e-5},
    {'params': model.layer2.parameters(), 'lr': 1e-5},
    {'params': model.layer3.parameters(), 'lr': 2e-5},
    {'params': model.layer4.parameters(), 'lr': 5e-5},
    {'params': model.fc.parameters(),     'lr': 1e-3},  # New head: highest LR
])

Data Augmentation for Small Datasets

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # ImageNet stats
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

train_dataset = ImageFolder("data/train", transform=train_transform)
val_dataset   = ImageFolder("data/val",   transform=val_transform)

Training Loop

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = model.to(device)
criterion = nn.CrossEntropyLoss()

def train_epoch(model, loader, optimizer, criterion):
    model.train()
    total_loss, correct = 0, 0
    for imgs, labels in loader:
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(imgs)
        loss    = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        correct    += (outputs.argmax(1) == labels).sum().item()
    return total_loss / len(loader), correct / len(loader.dataset)

Transfer Learning for NLP (HuggingFace)

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments

model_name = "distilbert-base-uncased"
tokenizer  = AutoTokenizer.from_pretrained(model_name)

# Fine-tune for 2-class sentiment analysis
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# Feature extraction: freeze all transformer layers, train only classifier
for name, param in model.named_parameters():
    if "classifier" not in name:
        param.requires_grad = False

# Full fine-tuning: just don't freeze anything
for param in model.parameters():
    param.requires_grad = True

Choosing the Right Pretrained Model

For image tasks in 2026: EfficientNetV2 and ConvNeXt are the best CNN options; Vision Transformer (ViT) variants offer state-of-the-art accuracy when you have enough data. For NLP: DistilBERT for speed, BERT-base for balance, RoBERTa-large for maximum accuracy. For lightweight mobile deployment: MobileNetV3 (vision) or DistilBERT (NLP). Always start with a smaller variant and scale up only if accuracy is insufficient.

Conclusion

Transfer learning is the single highest-leverage technique in deep learning. It collapses weeks of training into minutes and makes state-of-the-art results accessible with small labelled datasets. Always start with a pretrained model before considering training from scratch. With PyTorch and HuggingFace, the entire fine-tuning workflow is just a few dozen lines of code.

Leave feedback about this

  • Rating

Durgesh Kekare
Durgesh Kekarehttps://www.dataexpertise.in
Durgesh Kekare is a data science educator and founder of DataExpertise.in. With expertise in Python, machine learning, and analytics, he helps 10,000+ learners break into data careers.

Latest Posts

List of Categories