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

> Master transfer learning in PyTorch and TensorFlow. Reuse pre-trained networks to build accurate models fast with less data. Learn how to implement this powerful technique for your AI projects.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: tutorial
- Published: 2026-08-09

---

**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`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py) scan the directory to remove corrupt files before training begins.

```python
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`](https://github.com/microsoft/AI-For-Beginners/blob/main/tfcv.py). The same normalization statistics apply to ensure the VGG-16 backbone receives inputs in the expected range.

```python
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.

```python
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`.

```python
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.

```python
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.

```python
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`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py), which implements an `NLLLoss` criterion after `log_softmax` and tracks accuracy across epochs.

```python
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.

```python
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**

```python

# 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**

```python
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

```python
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

```python
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`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py), and [`tfcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/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`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py) and [`tfcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/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.