How to Implement Transfer Learning with Pre-Trained Networks in PyTorch and TensorFlow

Transfer learning lets you reuse a pre-trained model's feature-extraction layers and train only a new classification head on your specific dataset, achieving high accuracy with minimal data and training time.

The microsoft/AI-For-Beginners curriculum demonstrates this workflow through a cats-vs-dogs image classification example using the VGG-16 architecture. In this guide, you will learn how to implement transfer learning with pre-trained networks in PyTorch and TensorFlow by freezing backbone layers, replacing classifiers, and applying proper ImageNet normalization.

Preparing the Dataset with ImageNet Normalization

Both frameworks require consistent preprocessing to match the distribution of the original pre-training data. The notebooks in lessons/4-ComputerVision/08-TransferLearning/ use 224×224 pixel inputs with mean subtraction and standard deviation scaling.

PyTorch Data Pipeline

In TransferLearningPyTorch.ipynb, the torchvision.datasets.ImageFolder loads images from disk while transforms.Compose handles augmentation and normalization. The helper functions check_image_dir and check_image from pytorchcv.py scan the directory to remove corrupt files before training begins.

import torchvision.transforms as T
from torchvision.datasets import ImageFolder
from pytorchcv import check_image_dir

# Verify images are valid before loading

check_image_dir('data/PetImages')

transform = T.Compose([
    T.Resize(256),
    T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
])

dataset = ImageFolder('data/PetImages', transform=transform)

TensorFlow Data Pipeline

The TensorFlow implementation in TransferLearningTF.ipynb uses keras.preprocessing.image_dataset_from_directory or the custom load_cats_dogs_dataset helper from tfcv.py. The same normalization statistics apply to ensure the VGG-16 backbone receives inputs in the expected range.

import tensorflow as tf
from tfcv import check_image_dir, load_cats_dogs_dataset

# Clean corrupt files using shared utilities

check_image_dir('data/PetImages')

# Returns batched tf.data.Dataset objects

train_ds, val_ds = load_cats_dogs_dataset(batch_size=32)

Loading the Pre-Trained VGG-16 Backbone

Feature reuse relies on convolutional filters trained on ImageNet. You load these weights and immediately freeze them to prevent catastrophic forgetting during the initial training phase.

PyTorch Model Initialization

The PyTorch notebook loads torchvision.models.vgg16(pretrained=True) and sets requires_grad = False on all backbone parameters. This preserves the pre-trained edge and texture detectors while keeping memory usage lower during backpropagation.

import torch
import torchvision

vgg = torchvision.models.vgg16(pretrained=True)

# Freeze feature extraction layers

for p in vgg.features.parameters():
    p.requires_grad = False

TensorFlow Model Initialization

In TransferLearningTF.ipynb, tf.keras.applications.VGG16 instantiates the model with include_top=False to exclude the original 1,000-class ImageNet head. Setting base.trainable = False achieves the same effect as PyTorch's requires_grad.

from tensorflow import keras

base_model = keras.applications.VGG16(
    include_top=False,
    weights='imagenet',
    input_shape=(224, 224, 3)
)

base_model.trainable = False

Replacing the Classification Head

With the frozen backbone in place, you replace the final classification layers to match your specific number of output classes (two for cats vs. dogs).

PyTorch Classifier Replacement

The original VGG-16 classifier expects 25,088 input features. The code in TransferLearningPyTorch.ipynb substitutes a new nn.Sequential containing a custom linear layer mapping to two outputs.

from torch import nn

vgg.classifier = nn.Sequential(
    nn.Linear(25088, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 2)  # Binary classification

)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
vgg = vgg.to(device)

TensorFlow Head Construction

TensorFlow uses GlobalAveragePooling2D to reduce spatial dimensions before the dense layers. This approach reduces parameters compared to a flatten operation while maintaining spatial hierarchy information.

model = keras.Sequential([
    base_model,
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(2, activation='softmax')
])

Training Strategies: Freezing and Fine-Tuning

Transfer learning typically follows a two-phase approach: train the head first, then optionally unfreeze deeper layers for domain-specific refinement.

Phase 1: Training the Frozen Backbone

PyTorch uses the custom train() function from pytorchcv.py, which implements an NLLLoss criterion after log_softmax and tracks accuracy across epochs.

from torch.utils.data import DataLoader

train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
test_loader = DataLoader(test_set, batch_size=32)

history = train(vgg, train_loader, test_loader, epochs=5, lr=1e-3)

TensorFlow compiles the model with sparse categorical crossentropy and trains only the top layers because the base remains frozen.

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

history = model.fit(train_ds, validation_data=val_ds, epochs=5)

Phase 2: Optional Fine-Tuning

Unfreezing allows the backbone to adapt to domain-specific features. You must use a significantly lower learning rate (e.g., 1e-4 or 1e-5) to avoid destroying the pre-trained weights.

PyTorch Fine-Tuning

# Unfreeze last convolutional block

for p in vgg.features[24:].parameters():
    p.requires_grad = True

# Continue training with reduced learning rate

history = train(vgg, train_loader, test_loader, epochs=3, lr=1e-4)

TensorFlow Fine-Tuning

base_model.trainable = True

model.compile(
    optimizer=keras.optimizers.Adam(1e-4),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

model.fit(train_ds, validation_data=val_ds, epochs=3)

Complete Working Examples

These self-contained snippets replicate the core workflow from lessons/4-ComputerVision/08-TransferLearning/ without the helper modules, showing the essential API calls for each framework.

Full PyTorch Implementation

import torch
import torchvision
import torchvision.transforms as T
from torch import nn
from torch.utils.data import DataLoader, random_split

# Data preparation

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

full_ds = torchvision.datasets.ImageFolder('data/PetImages', transform)
train_ds, test_ds = random_split(full_ds, [20000, len(full_ds)-20000])

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=32)

# Model setup

vgg = torchvision.models.vgg16(pretrained=True)
for p in vgg.features.parameters():
    p.requires_grad = False

vgg.classifier = nn.Sequential(
    nn.Linear(25088, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 2)
)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
vgg = vgg.to(device)

# Training

optimizer = torch.optim.Adam(vgg.classifier.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(5):
    for imgs, labels in train_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = vgg(imgs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1} complete")

Full TensorFlow Implementation

import tensorflow as tf
from tensorflow import keras

# Data preparation

train_ds = keras.preprocessing.image_dataset_from_directory(
    'data/PetImages',
    validation_split=0.2,
    subset='training',
    seed=123,
    image_size=(224, 224),
    batch_size=32
)

val_ds = keras.preprocessing.image_dataset_from_directory(
    'data/PetImages',
    validation_split=0.2,
    subset='validation',
    seed=123,
    image_size=(224, 224),
    batch_size=32
)

# Model setup

base = keras.applications.VGG16(
    include_top=False,
    weights='imagenet',
    input_shape=(224, 224, 3)
)
base.trainable = False

model = keras.Sequential([
    base,
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(2, activation='softmax')
])

# Training

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_ds, validation_data=val_ds, epochs=5)

Summary

• Transfer learning accelerates development by reusing VGG-16 features trained on ImageNet, requiring only a small dataset to achieve high accuracy on new tasks like cats-vs-dogs classification.
• Normalization consistency is critical—use mean [0.485, 0.456, 0.406] and std [0.229, 0.224, 0.225] to match the pre-training distribution.
• Layer freezing prevents overfitting and reduces training time; freeze the backbone initially by setting requires_grad=False in PyTorch or trainable=False in TensorFlow.
• Head replacement involves substituting the final classification layer with a custom output sized to your number of classes (two in the provided examples).
• Fine-tuning optionally unfreezes deeper convolutional layers after initial convergence, using a smaller learning rate (1e-4) to refine domain-specific features without catastrophic forgetting.
• Source files for reference include TransferLearningPyTorch.ipynb, TransferLearningTF.ipynb, pytorchcv.py, and tfcv.py in the lessons/4-ComputerVision/08-TransferLearning/ directory of the microsoft/AI-For-Beginners repository.

Frequently Asked Questions

What is the primary benefit of freezing layers during transfer learning?

Freezing the pre-trained backbone preserves the generic visual features (edges, textures, shapes) learned from ImageNet while preventing overfitting on small datasets. This approach drastically reduces training time and computational cost because gradients only flow through the new classification head, not the entire deep network.

Which pre-trained architecture does the AI-For-Beginners curriculum use for demonstration?

The curriculum uses VGG-16 as the demonstration backbone in both TransferLearningPyTorch.ipynb and TransferLearningTF.ipynb. This architecture provides clear visualization of layer freezing and replacement patterns, though the same workflow applies to ResNet, MobileNet, or other models available in torchvision and tf.keras.applications.

How does the repository handle corrupt image files in the dataset?

The helper modules pytorchcv.py and tfcv.py provide the check_image_dir() and check_image() functions, which iterate through the dataset before training and remove files that cannot be decoded. This prevents runtime errors during the training loop when ImageFolder or image_dataset_from_directory encounters malformed JPEGs or truncated files.

When should you fine-tune the pre-trained backbone versus only training the new head?

You should train only the new head first to establish a baseline and prevent destroying pre-trained features. If validation accuracy plateaus or your target domain differs significantly from ImageNet (e.g., medical imaging), unfreeze the top few convolutional layers (or the entire backbone) and continue training with a 10×–100× smaller learning rate (e.g., 1e-4 or 1e-5) to adapt the filters to your specific data distribution.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}