Fine-tuning is a powerful technique in deep learning where a pre-trained model is adapted to a new, often related, task. This approach leverages the knowledge the model has already learned from a large dataset, reducing the amount of data and training time required for the new task. In this section, we will cover the following:

  1. Understanding Fine-Tuning
  2. Steps to Fine-Tune a CNN
  3. Practical Example: Fine-Tuning a Pre-trained Model
  4. Exercises

Understanding Fine-Tuning

Fine-tuning involves taking a pre-trained model, such as those available in PyTorch's torchvision.models module, and adjusting it to perform well on a new dataset. The key concepts include:

  • Transfer Learning: Using a model trained on one task and applying it to a different but related task.
  • Feature Extraction: Using the pre-trained model's learned features as input to a new classifier.
  • Fine-Tuning: Unfreezing some layers of the pre-trained model and jointly training both the pre-trained layers and the new layers.

Steps to Fine-Tune a CNN

  1. Load a Pre-trained Model: Choose a model pre-trained on a large dataset like ImageNet.
  2. Modify the Final Layer: Replace the final classification layer to match the number of classes in your new dataset.
  3. Freeze Layers: Optionally, freeze some of the earlier layers to retain their learned features.
  4. Train the Model: Train the modified model on your new dataset.

Practical Example: Fine-Tuning a Pre-trained Model

Let's walk through a practical example of fine-tuning a ResNet model for a new 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: Modify the Final Layer

import torch.nn as nn

# Modify the final layer to match the number of classes in the new dataset
num_classes = 10  # Example: 10 classes in the new dataset
model.fc = nn.Linear(model.fc.in_features, num_classes)

Step 3: Freeze Layers (Optional)

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

# Unfreeze the final layer
for param in model.fc.parameters():
    param.requires_grad = True

Step 4: Train the Model

import torch.optim as optim
from torchvision import datasets, transforms

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

# Load your new dataset
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
])

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

# Training loop
model.train()
for epoch in range(10):  # Example: 10 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)}')

Exercises

Exercise 1: Fine-Tune a Different Pre-trained Model

  1. Choose a different pre-trained model from torchvision.models (e.g., VGG16).
  2. Modify the final layer to match the number of classes in a new dataset.
  3. Train the model on the new dataset.

Exercise 2: Experiment with Freezing Layers

  1. Fine-tune a pre-trained model with different layers frozen.
  2. Compare the performance of the model with different configurations of frozen layers.

Solutions

Solution to Exercise 1

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

# Modify the final layer
model.classifier[6] = nn.Linear(model.classifier[6].in_features, num_classes)

# Freeze layers (optional)
for param in model.features.parameters():
    param.requires_grad = False

# Train the model (similar to the example above)

Solution to Exercise 2

# Experiment with freezing different layers
# Example: Freeze only the first few layers
for name, param in model.named_parameters():
    if 'features.0' in name or 'features.1' in name:
        param.requires_grad = False
    else:
        param.requires_grad = True

# Train the model (similar to the example above)

Conclusion

In this section, we covered the concept of fine-tuning CNNs, the steps involved, and provided a practical example. Fine-tuning allows you to leverage pre-trained models to achieve good performance on new tasks with less data and training time. By experimenting with different pre-trained models and configurations, you can optimize the performance for your specific use case.

© Copyright 2024. All rights reserved