# How to Implement Transfer Learning with Pre-Trained Models in AI For Beginners Computer Vision

> Learn transfer learning with pre-trained models in AI For Beginners computer vision. Repurpose ImageNet models for custom image classification easily. Freeze backbone, add new head, and fine-tune.

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

---

**Transfer learning in the AI For Beginners curriculum lets you repurpose ImageNet-trained models like MobileNetV2 or ResNet-18 for custom image classification by freezing the backbone, adding a new classification head, and optionally fine-tuning top layers.**

The computer vision unit in Microsoft's [AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners) repository provides complete, runnable implementations in both TensorFlow/Keras and PyTorch. These lessons demonstrate how to solve the classic "cats vs. dogs" classification problem using pre-trained weights rather than training from scratch.

## Why Transfer Learning Works for Small Datasets

Training deep neural networks requires millions of labeled images. Most practitioners don't have this data. **Transfer learning** bridges this gap by reusing features learned from large datasets (typically ImageNet's 1.2 million images) and adapting them to your specific task.

The AI For Beginners approach follows a proven two-phase pattern:

1. **Feature extraction** — freeze the pre-trained backbone and train only the new classification head
2. **Fine-tuning** — unfreeze select layers and continue training with reduced learning rates

This method achieves high accuracy with as few as 2,000 training images.

## TensorFlow/Keras Implementation

The file `lessons/4-ComputerVision/08-TransferLearning/TransferLearningTF.ipynb` demonstrates transfer learning using Keras Applications.

### Loading the Pre-Trained Backbone

```python
import tensorflow as tf
from tensorflow import keras
from tfcv import load_cats_dogs_dataset, plot_results

# Load dataset

ds_train, ds_test = load_cats_dogs_dataset(batch_size=64)

# Load MobileNetV2 pre-trained on ImageNet

base = keras.applications.MobileNetV2(
    weights='imagenet',
    include_top=False,           # exclude original classification head

    input_shape=(224, 224, 3)
)

```

Setting `include_top=False` removes ImageNet's 1000-class output layer, allowing you to attach your own classifier.

### Freezing and Building the Model

```python

# Freeze all backbone layers

base.trainable = False

# Build new classification head

inputs = keras.Input(shape=(224, 224, 3))
x = keras.applications.mobilenet_v2.preprocess_input(inputs)
x = base(x, training=False)    # run in inference mode (BatchNorm)

x = keras.layers.GlobalAveragePooling2D()(x)
outputs = keras.layers.Dense(2, activation='softmax')(x)

model = keras.Model(inputs, outputs)

```

The `GlobalAveragePooling2D` layer reduces spatial dimensions from `(7, 7, 1280)` to `(1280)`, feeding directly into the 2-class output for cats vs. dogs.

### Training and Fine-Tuning

```python

# Phase 1: Train only the head

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

# Phase 2: Fine-tune top layers

base.trainable = True
model.compile(
    optimizer=keras.optimizers.Adam(1e-5),  # 100x smaller learning rate

    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)
fine_history = model.fit(
    ds_train,
    validation_data=ds_test,
    epochs=3
)

```

The helper function `plot_results(history)` in [`tfcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/tfcv.py) generates visualization of training curves.

## PyTorch Implementation

The equivalent PyTorch workflow appears in `lessons/4-ComputerVision/08-TransferLearning/TransferLearningPyTorch.ipynb`, using `torchvision` models.

### Model Setup and Layer Freezing

```python
import torch
import torchvision
from pytorchcv import load_cats_dogs_dataset, train, validate, plot_results

# Load dataset with custom loaders

dataset, train_loader, test_loader = load_cats_dogs_dataset()

# Load pre-trained ResNet-18

base = torchvision.models.resnet18(pretrained=True)

# Freeze all parameters

for param in base.parameters():
    param.requires_grad = False

# Replace final FC layer for binary classification

num_ftrs = base.fc.in_features
base.fc = torch.nn.Linear(num_ftrs, 2)

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

```

### Training Loop and Selective Unfreezing

```python

# Phase 1: Train classification head only

hist = train(base, train_loader, test_loader, epochs=5, lr=0.001)

# Phase 2: Unfreeze last residual block (layer4) for fine-tuning

for name, param in base.named_parameters():
    if "layer4" in name:
        param.requires_grad = True

# Re-create optimizer with only trainable parameters

optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, base.parameters()),
    lr=1e-5
)

fine_hist = train(base, train_loader, test_loader,
                  optimizer=optimizer, epochs=3, lr=1e-5)

```

The `train()` and `validate()` functions in [`pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py) implement standard training loops with progress tracking and metric logging.

## Key Architectural Differences: TensorFlow vs. PyTorch

| Aspect | TensorFlow/Keras | PyTorch |
|--------|----------------|---------|
| **Pre-trained model loading** | `keras.applications.MobileNetV2(weights='imagenet', include_top=False)` | `torchvision.models.resnet18(pretrained=True)` |
| **Freeze mechanism** | `base.trainable = False` | `param.requires_grad = False` for each parameter |
| **Classification head** | `GlobalAveragePooling2D()` → `Dense(2, activation='softmax')` | `AdaptiveAvgPool2d((1,1))` → replace `fc` layer |
| **Training orchestration** | `model.fit()` with callbacks | Custom `train()` function in [`pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py) |
| **Fine-tuning control** | Unfreeze entire backbone: `base.trainable = True` | Selective unfreeze by layer name: `"layer4" in name` |

## Helper Functions and Utilities

The repository includes utility modules that abstract boilerplate code:

- **[`tfcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/tfcv.py)** — dataset loading (`load_cats_dogs_dataset`), image preprocessing, and visualization (`plot_results`)
- **[`pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py)** — PyTorch dataset wrappers, training loops (`train`, `validate`), and metric plotting

These files let you focus on model architecture rather than data pipeline implementation.

## Fine-Tuning Strategies

Both implementations demonstrate progressive unfreezing:

| Strategy | When to Use | Learning Rate |
|----------|-------------|---------------|
| **Head-only training** | Very small datasets (< 1,000 images) | Standard (1e-3) |
| **Top-block fine-tuning** | Medium datasets, similar to ImageNet | Reduced (1e-4 to 1e-5) |
| **Full fine-tuning** | Large datasets, very different from ImageNet | Very low (1e-5 to 1e-6) |

The AI For Beginners notebooks use **top-block fine-tuning** as the default balanced approach.

## Adapting to Your Own Dataset

Replace the backbone and adjust class count:

**TensorFlow:**

```python

# Swap MobileNetV2 for EfficientNetB0

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

# Change output classes

outputs = keras.layers.Dense(NUM_YOUR_CLASSES, activation='softmax')(x)

```

**PyTorch:**

```python

# Swap ResNet-18 for DenseNet-121

base = torchvision.models.densenet121(pretrained=True)

# Replace classifier

base.classifier = torch.nn.Linear(base.classifier.in_features, NUM_YOUR_CLASSES)

```

## Summary

- **Transfer learning** in AI For Beginners uses a consistent two-phase approach across TensorFlow and PyTorch implementations
- **Freeze the backbone** initially to preserve ImageNet features while training a new classification head
- **Fine-tune selectively** with reduced learning rates to adapt deeper layers to your specific data
- **Key files**: `TransferLearningTF.ipynb`, `TransferLearningPyTorch.ipynb`, [`tfcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/tfcv.py), and [`pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py) in `lessons/4-ComputerVision/08-TransferLearning/`

## Frequently Asked Questions

### What pre-trained models does AI For Beginners recommend for beginners?

The official notebooks use **MobileNetV2** (TensorFlow) and **ResNet-18** (PyTorch) because they offer strong accuracy with moderate computational requirements. MobileNetV2 runs efficiently on CPU, making it ideal for educational environments without GPU access. Both models achieve >95% accuracy on the cats vs. dogs dataset after fine-tuning.

### How does freezing layers work in practice?

**Freezing** prevents gradient updates from modifying pre-trained weights. In Keras, `base.trainable = False` sets all layer `trainable` attributes to False. In PyTorch, you set `requires_grad=False` on each parameter tensor. This ensures the backbone acts as a fixed feature extractor during initial training, preventing overfitting on small datasets.

### When should I fine-tune versus just training the head?

Fine-tune when your **dataset size exceeds ~1,000 images per class** or when your target images differ significantly from ImageNet (medical imaging, satellite photos, industrial defects). If your data closely resembles natural photographs and you have limited samples, head-only training often suffices according to the patterns in [`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).

### Can I use these patterns for multi-label classification?

Yes. Replace the final activation: use **sigmoid** instead of **softmax** in TensorFlow, or `nn.BCEWithLogitsLoss` in PyTorch. Change the output dimension to match your number of labels. The backbone feature extraction remains identical; only the classification head and loss function require modification.