From Pretraining to Fine-Tuning: Strategies for Adapting Models to New Tasks — Part2

Introduction: Unlocking the Power of Pretrained Models

In the previous part, we explored how pretrained models can be saved for reuse and efficient deployment. Now, let’s shift our focus to the practical ways these models can be used for downstream tasks. Pretrained models act as a versatile foundation, allowing you to adapt them to a wide variety of applications with minimal training effort.

In this part, we’ll explore:

1. Fine-Tuning: Updating all layers of the model for task-specific optimization.
2. Linear Probing: Using the pretrained model as a feature extractor while training a lightweight classifier.
3. Feature Extraction: Leveraging embeddings or intermediate representations from pretrained models.
4. Task-Specific Adaptation: Modifying architectures for specific domains (e.g., object detection, sequence tagging).

By the end, you’ll understand how to harness pretrained models effectively for real-world tasks.

1. Fine-Tuning Pretrained Models

Fine-tuning is one of the most effective methods for leveraging pretrained models in downstream tasks. It involves updating all layers of the pretrained model to adapt it to a specific task. This approach is particularly useful when your downstream dataset is large or closely related to the original dataset used for pretraining.

Steps for Fine-Tuning

1. Load the Pretrained Model Start by loading a pretrained model, such as ResNet for image tasks or BERT for NLP tasks.
2. Replace the Output Layer Modify the final layer to match the number of classes or outputs for your specific task.
3. Freeze or Unfreeze Layers

• You can freeze some layers (especially earlier ones) to retain their pretrained knowledge.
• Fine-tune the entire model by unfreezing all layers.

Train with a smaller learning rate to avoid overwriting the pretrained weights.

Code Example: Fine-Tuning ResNet for Image Classification

import torch
import torchvision.models as models
import torch.nn as nn

# Load a pretrained ResNet model
model = models.resnet50(pretrained=True)

# Modify the final fully connected layer for a specific task (e.g., 10 classes)
num_classes = 10
model.fc = nn.Linear(model.fc.in_features, num_classes)

# Optionally freeze earlier layers
for param in model.parameters():
    param.requires_grad = False

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

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

# Training loop (simplified)
for epoch in range(5):
    for inputs, labels in train_loader:  # Assuming train_loader is defined
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

Code Example: Fine-Tuning BERT for Text Classification

from transformers import BertForSequenceClassification, AdamW

# Load a pretrained BERT model
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

# Define optimizer and learning rate scheduler
optimizer = AdamW(model.parameters(), lr=2e-5)

# Training loop (simplified)
for epoch in range(3):
    for batch in train_dataloader:  # Assuming train_dataloader is defined
        inputs = {key: val.to(device) for key, val in batch.items()}
        outputs = model(**inputs)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

Best Practices for Fine-Tuning

1. Start with a Low Learning Rate: Prevents overwriting pretrained knowledge.
2. Use Gradual Unfreezing: Unfreeze layers progressively to balance between retaining pretrained knowledge and adapting to new data.
3. Monitor Overfitting: Use techniques like dropout, regularization, and early stopping if working with a small dataset.

When to Use Fine-Tuning

• Your downstream task is similar to the pretraining task.
• You have a sufficiently large dataset for fine-tuning.
• You need high task-specific performance.

2. Linear Probing

Linear probing is a lightweight and computationally efficient method to leverage pretrained models. Unlike fine-tuning, linear probing involves freezing the pretrained model’s layers and training only a simple classifier on top. This method is particularly effective when you have limited data or when the downstream task is different from the pretraining task.

Steps for Linear Probing

1. Load the Pretrained Model Use a pretrained model and freeze all its layers to retain the knowledge acquired during pretraining.
2. Add a Classifier Attach a simple linear classifier (e.g., fully connected layer) on top of the pretrained model.
3. Train the Classifier Train only the newly added classifier using the frozen features of the pretrained model.

Code Example: Linear Probing with ResNet

import torch
import torchvision.models as models
import torch.nn as nn

# Load a pretrained ResNet model
model = models.resnet50(pretrained=True)

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

# Replace the final fully connected layer with a new classifier
num_classes = 10
model.fc = nn.Linear(model.fc.in_features, num_classes)

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

# Training loop (simplified)
for epoch in range(5):
    for inputs, labels in train_loader:  # Assuming train_loader is defined
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

Code Example: Linear Probing with BERT

from transformers import BertModel, BertTokenizer
import torch.nn as nn
import torch

# Load a pretrained BERT model
bert = BertModel.from_pretrained("bert-base-uncased")

# Freeze all layers
for param in bert.parameters():
    param.requires_grad = False

# Define a custom classifier on top of the frozen BERT model
class BertClassifier(nn.Module):
    def __init__(self, bert, num_classes):
        super(BertClassifier, self).__init__()
        self.bert = bert
        self.classifier = nn.Linear(bert.config.hidden_size, num_classes)
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        cls_output = outputs.pooler_output
        return self.classifier(cls_output)

# Instantiate and train the classifier
num_classes = 2
model = BertClassifier(bert, num_classes)

# Training loop similar to fine-tuning but only updates the classifier
optimizer = torch.optim.Adam(model.classifier.parameters(), lr=0.001)

Best Practices for Linear Probing

1. Use Pretrained Features Efficiently: Ensure the features from the pretrained model are well-suited for the downstream task.
2. Evaluate the Classifier: Compare performance across different classifier architectures (e.g., linear layer, MLP).
3. Start Simple: Linear probing often works surprisingly well with just one or two layers on top.

When to Use Linear Probing

• Your dataset is small, and full fine-tuning might lead to overfitting.
• The downstream task is substantially different from the pretraining task.
• You need a computationally efficient solution.

3. Feature Extraction

Feature extraction is another powerful way to use pretrained models for downstream tasks. Instead of training the entire model or adding new layers, you use the pretrained model to generate embeddings or intermediate representations from your data. These extracted features can then be used as input for simpler models like support vector machines (SVMs), random forests, or even fully connected networks.

This approach is especially useful when:

• You have limited computational resources.
• The downstream task is significantly different from the pretraining task.
• You need interpretable features.

Steps for Feature Extraction

1. Load the Pretrained Model Use a pretrained model like ResNet or BERT.
2. Extract Features Pass your data through the model and collect the outputs from intermediate layers.
3. Use the Features for New Tasks Train a lightweight model or directly use the features for analysis, clustering, or other downstream tasks.

Code Example: Feature Extraction with ResNet

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

# Load a pretrained ResNet model
model = models.resnet50(pretrained=True)

# Remove the final classification layer
feature_extractor = torch.nn.Sequential(*list(model.children())[:-1])
feature_extractor.eval()  # Set to evaluation mode

# Example input image
image = Image.open("example.jpg")
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Extract features
input_tensor = transform(image).unsqueeze(0)  # Add batch dimension
features = feature_extractor(input_tensor).squeeze().detach().numpy()
print("Extracted Features Shape:", features.shape)

Code Example: Feature Extraction with BERT

from transformers import BertModel, BertTokenizer
import torch

# Load a pretrained BERT model
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")

# Example input text
text = "Deep learning models are powerful."
inputs = tokenizer(text, return_tensors="pt")

# Extract features
with torch.no_grad():
    outputs = model(**inputs)
    hidden_states = outputs.last_hidden_state  # [batch_size, seq_length, hidden_size]

# Use the CLS token embedding for classification or other tasks
cls_embedding = hidden_states[:, 0, :].squeeze().numpy()
print("Extracted CLS Embedding Shape:", cls_embedding.shape)

Best Practices for Feature Extraction

Select the Right Layer:

• Use early layers for general-purpose features.
• Use later layers for task-specific features.

Preprocess Consistently: Ensure your input data is preprocessed the same way as the data used during pretraining.

Combine Features: Experiment with combining features from multiple layers to capture both low-level and high-level information.

When to Use Feature Extraction

• You want to reduce the complexity of the downstream task by using pretrained features.
• Your computational resources are limited, making full fine-tuning infeasible.
• You’re working on unsupervised or semi-supervised tasks like clustering or dimensionality reduction.

4. Task-Specific Adaptation

Task-specific adaptation involves modifying pretrained models to suit specialized tasks or domains. Unlike fine-tuning or feature extraction, this method requires structural changes to the model, such as adding new components or modifying existing ones. This approach is commonly used for tasks like object detection, sequence labeling, or multi-modal learning, where the default architecture of the pretrained model may not directly align with the target task.

Steps for Task-Specific Adaptation

1. Analyze the Task Requirements Understand the input-output format and specific requirements of the downstream task (e.g., bounding boxes for object detection or BIO tags for sequence labeling).
2. Modify the Model Architecture Extend the pretrained model by adding task-specific components, such as:

• Detection heads for object detection.
• CRF layers for sequence labeling.
• Fusion layers for multi-modal tasks.

Train the Adapted Model Fine-tune the modified architecture using task-specific datasets.

Example 1: Adapting a Vision Model for Object Detection

Object detection involves detecting and classifying objects within an image. Pretrained image classification models can be adapted by adding detection heads.

from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# Load a pretrained Faster R-CNN model
model = fasterrcnn_resnet50_fpn(pretrained=True)

# Modify the box predictor to match the number of classes in the new dataset
num_classes = 5  # Including the background class
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# Train the modified model (training loop not shown for brevity)

Example 2: Adapting BERT for Sequence Labeling

Sequence labeling tasks, such as named entity recognition (NER), require token-level predictions. BERT can be adapted for this task by adding a token classification head.

from transformers import BertForTokenClassification

# Load a pretrained BERT model with a token classification head
model = BertForTokenClassification.from_pretrained("bert-base-uncased", num_labels=9)
# Train the model (training loop not shown for brevity)

Example 3: Multi-Modal Learning with CLIP

Multi-modal tasks, like image-text retrieval, require combining representations from different modalities. CLIP, a multi-modal pretrained model, can be adapted to tasks like image-captioning or visual question answering.

from transformers import CLIPProcessor, CLIPModel

# Load a pretrained CLIP model
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Example input: Image and text pair
inputs = processor(text=["A cat sitting on a mat"], images=["cat_image.jpg"], return_tensors="pt", padding=True)

# Forward pass to get embeddings
outputs = model(**inputs)
image_embeds = outputs.image_embeds
text_embeds = outputs.text_embeds

Best Practices for Task-Specific Adaptation

1. Understand Model Limitations: Ensure the pretrained model is suitable for the task (e.g., spatial information for vision tasks).
2. Experiment with Modifications: Test different architectures and task-specific heads to find the optimal setup.
3. Leverage Specialized Frameworks: Use libraries like Hugging Face for NLP tasks or Detectron2 for vision tasks to streamline adaptation.

When to Use Task-Specific Adaptation

• The task output requires a specific structure not supported by the default model architecture.
• You’re working on complex tasks like object detection, sequence labeling, or multi-modal learning.
• Pretrained models provide a solid base but need additional components to meet task requirements.

Conclusion

Pretrained models have transformed the landscape of deep learning, enabling developers and researchers to achieve state-of-the-art results with reduced time and resources. In this two-part blog series, we explored:

1. Part 1: How to save pretrained models effectively for reuse, covering techniques like saving entire models, weights, encoders, and checkpoints.
2. Part 2: Various strategies to leverage pretrained models for downstream tasks, including fine-tuning, linear probing, feature extraction, and task-specific adaptations.

These approaches empower you to unlock the full potential of pretrained models, whether you’re working with massive datasets or tackling resource-constrained environments. Fine-tuning provides task-specific optimization, while linear probing and feature extraction offer lightweight alternatives. For more complex tasks, task-specific adaptations open doors to customized solutions.

The key to success lies in understanding your task requirements, choosing the right strategy, and experimenting with these techniques to find the optimal setup for your use case. Pretrained models are not just tools but catalysts for innovation, enabling you to build efficient, scalable, and impactful AI solutions.

Stay tuned for future content where we’ll delve deeper into advanced techniques for model optimization, deployment, and real-world applications!