# Implement Transfer Learning in PyTorch vs TensorFlow: A Complete Beginner's Guide

> Learn to implement transfer learning in PyTorch vs TensorFlow with this beginner's guide. Discover how to reuse pretrained models for faster AI development.

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

---

**Transfer learning accelerates model training by reusing pretrained neural network weights as feature extractors, with PyTorch using `torchvision.models` and TensorFlow using `tf.keras.applications` to freeze backbones and replace classification heads.**

Transfer learning is a cornerstone technique in modern computer vision that allows practitioners to leverage pretrained models like ResNet and MobileNet for new classification tasks. The Microsoft AI-For-Beginners repository provides comprehensive, side-by-side implementations showing exactly how to implement transfer learning in PyTorch vs TensorFlow for educational workflows. These implementations demonstrate freezing convolutional backbones, replacing final layers, and fine-tuning on custom datasets like Cats vs Dogs.

## PyTorch Implementation: ResNet-18 Feature Extraction

The PyTorch implementation in [`lessons/4-ComputerVision/08-TransferLearning/pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/08-TransferLearning/pytorchcv.py) demonstrates a complete transfer learning pipeline using **ResNet-18** pretrained on ImageNet.

### Freezing the Convolutional Backbone

The code loads a pretrained model via `torchvision.models.resnet18(pretrained=True)` and immediately freezes all parameters to prevent gradient updates during initial training. This preserves the generic visual features learned from ImageNet.

```python
import torch
import torch.nn as nn
import torchvision

# Load pretrained backbone

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

# Freeze all layers

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

```

### Replacing the Classification Head

After freezing, the final fully-connected layer is replaced with a new `nn.Linear` layer matching the target dataset's class count. In the AI-For-Beginners curriculum, this adapts the model from 1000 ImageNet classes to binary classification (Cats vs Dogs) or 10 MNIST digits.

```python
num_classes = 2  # Cats vs Dogs

model.fc = nn.Linear(model.fc.in_features, num_classes)

```

### Training Utilities and Data Pipeline

The [`pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py) module provides helper functions including `load_mnist`, `load_cats_dogs_dataset`, and `common_transform` to standardize preprocessing. The training loop uses `train_epoch` and `validate` utilities that implement the standard PyTorch pattern: forward pass, loss calculation, backward propagation, and optimizer steps.

```python
from torchvision import transforms, datasets

# ImageNet standard preprocessing

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

train_ds = datasets.ImageFolder('data/PetImages', transform=transform)
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)

```

## TensorFlow Implementation: MobileNetV2 Fine-Tuning

The TensorFlow implementation in `lessons/4-ComputerVision/08-TransferLearning/TransferLearningTF.ipynb` uses **MobileNetV2** through the Keras Applications API, offering a higher-level abstraction while maintaining the same core concepts.

### Loading and Freezing the Base Model

The implementation loads `tf.keras.applications.MobileNetV2` with `include_top=False` to exclude the original classification head and `weights='imagenet'` to load pretrained parameters. Setting `base.trainable = False` freezes the convolutional backbone.

```python
import tensorflow as tf
from tensorflow.keras import layers, models

# Load pretrained feature extractor

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

```

### Building the Custom Classification Head

Rather than modifying an existing layer, TensorFlow uses the functional API to stack new layers on top of the frozen base. A `GlobalAveragePooling2D` layer reduces spatial dimensions before a `Dense` layer with softmax activation performs the final classification.

```python
inputs = tf.keras.Input(shape=(224, 224, 3))
x = base(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(2, activation='softmax')(x)  # 2 classes

model = models.Model(inputs, outputs)

```

### Optional Fine-Tuning Stage

Unlike the basic PyTorch example, the TensorFlow notebook demonstrates **fine-tuning** by unfreezing deeper layers after initial training. This involves setting `base_model.trainable = True` and recompiling with a lower learning rate to prevent catastrophic forgetting.

```python

# Unfreeze for fine-tuning

base.trainable = True
model.compile(optimizer=tf.keras.optimizers.Adam(1e-5),  # Lower LR

              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

```

## Complete End-to-End Examples

Below are self-contained snippets extracted from the AI-For-Beginners curriculum that you can run immediately after installing dependencies listed in the repository's [`requirements.txt`](https://github.com/microsoft/AI-For-Beginners/blob/main/requirements.txt).

### PyTorch: Full Training Loop

```python
import torch
import torch.nn as nn
import torchvision
from torchvision import transforms, datasets

# 1. Load pretrained model and freeze backbone

model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

# 2. Replace classifier head for binary classification

num_classes = 2
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 3. Prepare data with ImageNet normalization

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])
train_ds = datasets.ImageFolder('data/PetImages', transform=transform)
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)

# 4. Train only the new head

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

for epoch in range(5):
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

```

### TensorFlow: Dataset Pipeline and Training

```python
import tensorflow as tf
from tensorflow.keras import layers, models

# 1. Load frozen pretrained base

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

# 2. Build new head with pooling

inputs = tf.keras.Input(shape=(224, 224, 3))
x = base(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(2, activation='softmax')(x)
model = models.Model(inputs, outputs)

# 3. Load and preprocess images

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    'data/PetImages',
    image_size=(224, 224),
    batch_size=32
)

# 4. Compile and fit

model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-3),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)
model.fit(train_ds, epochs=5)

```

## Key Files in the Repository

| File Path | Framework | Description |
|-----------|-----------|-------------|
| [`lessons/4-ComputerVision/08-TransferLearning/pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/08-TransferLearning/pytorchcv.py) | PyTorch | Complete training utilities, data loaders, and visualization functions |
| `lessons/4-ComputerVision/08-TransferLearning/TransferLearningPyTorch.ipynb` | PyTorch | Interactive notebook with exploratory data analysis |
| `lessons/4-ComputerVision/08-TransferLearning/TransferLearningTF.ipynb` | TensorFlow | Full fine-tuning workflow with callbacks and validation |
| [`requirements.txt`](https://github.com/microsoft/AI-For-Beginners/blob/main/requirements.txt) | Both | Dependency specifications for TensorFlow 2.x and PyTorch 1.x |

## Summary

- **Freeze the backbone** to preserve ImageNet features using `requires_grad=False` (PyTorch) or `trainable=False` (TensorFlow).
- **Replace the head** with framework-specific layers (`nn.Linear` vs `layers.Dense`) matching your class count.
- **Use ImageNet normalization** consistently across both frameworks with mean `[0.485, 0.456, 0.406]` and std `[0.229, 0.224, 0.225]`.
- **Fine-tune optionally** by unfreezing deeper layers and reducing learning rate, as demonstrated in the TensorFlow notebook.
- **Reference [`pytorchcv.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/pytorchcv.py)** for reusable PyTorch training loops and `TransferLearningTF.ipynb` for Keras callbacks and data augmentation.

## Frequently Asked Questions

### How do I choose between PyTorch and TensorFlow for transfer learning?

**PyTorch** offers explicit control over gradient computation and model architecture, making it ideal for research and custom training loops. **TensorFlow/Keras** provides higher-level abstractions like `model.fit()` and built-in callbacks, streamlining production workflows. Both frameworks achieve identical accuracy when implementing the same architecture and preprocessing.

### Can I use different pretrained models than ResNet-18 or MobileNetV2?

Yes. In PyTorch, substitute `resnet18` with `torchvision.models.efficientnet_b0` or `vgg16`. In TensorFlow, replace `MobileNetV2` with `tf.keras.applications.ResNet50` or `EfficientNetB0`. Ensure you adjust the input preprocessing to match each model's expected normalization statistics.

### Why must I freeze layers during initial training?

Freezing prevents overfitting when your dataset is small (e.g., the Cats vs Dogs subset in AI-For-Beginners) and preserves generic low-level features (edges, textures) learned from millions of ImageNet images. This allows training to converge in minutes rather than hours while maintaining high accuracy.

### How do I implement fine-tuning after the initial frozen training?

In PyTorch, unfreeze specific layers by setting `param.requires_grad = True` and create a new optimizer with a lower learning rate for those parameters. In TensorFlow, set `base_model.trainable = True` and recompile the model with `optimizer=tf.keras.optimizers.Adam(1e-5)` before calling `model.fit()` again for additional epochs.