The Project | About Us | Contribute | Donations | License
HOME My courses Completed
◄ Previous
◄

Transfer Learning with Pre-trained Models

Next ►
►

Transfer learning is a powerful technique in deep learning where a model developed for a particular task is reused as the starting point for a model on a second task. This is particularly useful when you have limited data for the second task. In this section, we will explore how to leverage pre-trained models in PyTorch for transfer learning.

Key Concepts

  1. Pre-trained Models: Models that have been previously trained on large datasets, such as ImageNet.
  2. Feature Extraction: Using the convolutional base of a pre-trained model to extract features from new data.
  3. Fine-Tuning: Unfreezing some of the top layers of a frozen model base and jointly training both the newly added part of the model and the unfrozen layers.

Steps for Transfer Learning

  1. Load a Pre-trained Model: PyTorch provides several pre-trained models through the torchvision.models module.
  2. Freeze the Convolutional Base: Prevent the weights of the pre-trained model from being updated during training.
  3. Add Custom Layers: Add new layers that are specific to your task.
  4. Train the Model: Train the new layers while keeping the pre-trained layers frozen, or fine-tune the entire model.

Practical Example

Let's walk through a practical example of transfer learning using a pre-trained ResNet model for an image classification task.

Step 1: Load a Pre-trained Model

import torch
import torchvision.models as models

# Load a pre-trained ResNet model
model = models.resnet18(pretrained=True)

Step 2: Freeze the Convolutional Base

# Freeze all the parameters in the model
for param in model.parameters():
    param.requires_grad = False

Step 3: Add Custom Layers

import torch.nn as nn

# Replace the final fully connected layer
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10)  # Assuming we have 10 classes

Step 4: Train the Model

import torch.optim as optim

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

# Training loop
for epoch in range(10):  # Number of epochs
    running_loss = 0.0
    for inputs, labels in dataloader:  # Assuming dataloader is defined
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(dataloader)}")

Fine-Tuning the Model

If you want to fine-tune the entire model, you can unfreeze some of the layers:

# Unfreeze the last few layers
for param in model.layer4.parameters():
    param.requires_grad = True

# Update the optimizer to include all parameters
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

Practical Exercise

Exercise: Use a pre-trained VGG16 model to classify images from a custom dataset with 5 classes. Follow the steps outlined above to perform transfer learning.

Solution

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

# Load a pre-trained VGG16 model
model = models.vgg16(pretrained=True)

# Freeze all the parameters in the model
for param in model.parameters():
    param.requires_grad = False

# Replace the final classifier layer
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 5)  # Assuming we have 5 classes

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.classifier[6].parameters(), lr=0.001, momentum=0.9)

# Data loading and preprocessing
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

train_dataset = datasets.ImageFolder(root='path/to/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# Training loop
for epoch in range(10):  # Number of epochs
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}")

Summary

In this section, we covered the concept of transfer learning and how to implement it using pre-trained models in PyTorch. We walked through the steps of loading a pre-trained model, freezing its layers, adding custom layers, and training the model. We also provided a practical exercise to reinforce the learned concepts. Transfer learning is a powerful technique that can significantly reduce the training time and improve the performance of your models, especially when you have limited data.

◄ Previous
◄
Next ►
►
The Project | About Us | Contribute | Donations | License
© Copyright 2024. All rights reserved