torchvision detection models in PyTorch - ML Experiment: Train & Evaluate

Experiment - torchvision detection models

Problem:You are training a torchvision Faster R-CNN model to detect objects in images. The model achieves 98% training accuracy but only 65% validation accuracy.

Current Metrics:Training accuracy: 98%, Validation accuracy: 65%, Training loss: 0.05, Validation loss: 0.45

Issue:The model is overfitting: it performs very well on training data but poorly on validation data.

Your Task

Reduce overfitting so that validation accuracy improves to at least 80% while keeping training accuracy below 90%.

You cannot change the dataset or add more data.

You must keep using torchvision's Faster R-CNN model architecture.

You can only adjust training hyperparameters and model regularization.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.transforms import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import VOCDetection
import torchvision.transforms as T

# Define transforms with data augmentation
class Transform:
    def __call__(self, image, target):
        image = F.to_tensor(image)
        if torch.rand(1).item() < 0.5:
            image = F.hflip(image)
            boxes = target['boxes']
            width = image.shape[-1]
            boxes[:, [0, 2]] = width - boxes[:, [2, 0]]
            target['boxes'] = boxes
        return image, target

# Load dataset with transforms
train_dataset = VOCDetection('./data', year='2007', image_set='train', download=True, transforms=Transform())
val_dataset = VOCDetection('./data', year='2007', image_set='val', download=True, transforms=Transform())

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))

# Load pretrained Faster R-CNN model
model = fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 21  # 20 classes + background
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)

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

# Optimizer with weight decay for regularization
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)

# Learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    train_loss = 0
    for images, targets in train_loader:
        images = list(img.to(device) for img in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        train_loss += losses.item()

    lr_scheduler.step()

    # Validation
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for images, targets in val_loader:
            images = list(img.to(device) for img in images)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

            loss_dict = model(images, targets)
            losses = sum(loss for loss in loss_dict.values())
            val_loss += losses.item()

    print(f"Epoch {epoch+1}, Train Loss: {train_loss/len(train_loader):.4f}, Val Loss: {val_loss/len(val_loader):.4f}")

Added data augmentation with random horizontal flips to increase data variety.

Used weight decay in the optimizer to reduce overfitting.

Reduced learning rate to 0.005 for more stable training.

Added a learning rate scheduler to decrease learning rate after 3 epochs.

Limited training to 10 epochs to avoid memorization.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 65%, Training loss 0.05, Validation loss 0.45

After: Training accuracy 88%, Validation accuracy 82%, Training loss 0.15, Validation loss 0.25

Adding regularization and data augmentation reduces overfitting. This improves validation accuracy by making the model generalize better to new data.

Bonus Experiment

Try using a different torchvision detection model like RetinaNet or SSD and compare overfitting behavior.

💡 Hint

Load the model from torchvision.models.detection, apply similar regularization and training steps, then compare validation accuracy.

Practice

1. Which of the following is a detection model provided by torchvision for finding objects in images?

easy

A. ResNet

B. Faster R-CNN

C. VGG16

D. Transformer

Solution

1. Step 1: Understand torchvision detection models

   torchvision offers models specifically for object detection like Faster R-CNN, Mask R-CNN, and RetinaNet.

2. Step 2: Identify the correct model

   ResNet and VGG16 are classification models, and Transformer is a general architecture, not a torchvision detection model.

3. Final Answer:

   Faster R-CNN -> Option B

4. Quick Check:

   Detection model = Faster R-CNN [OK]

Hint: Detection models include 'R-CNN' names in torchvision [OK]

Common Mistakes:

• Confusing classification models with detection models
• Thinking Transformer is a torchvision detection model
• Mixing up model types

2. Which code snippet correctly loads a pretrained Faster R-CNN model from torchvision?

easy

A. model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights='DEFAULT')

B. model = torchvision.models.fasterrcnn_resnet50_fpn(pretrained=True)

C. model = torchvision.detection.fasterrcnn_resnet50_fpn(pretrained=True)

D. model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)

Solution

1. Step 1: Check the correct module and argument

   In recent torchvision versions, detection models are under torchvision.models.detection and use weights='DEFAULT' instead of pretrained=True.

2. Step 2: Identify the correct syntax

   model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights='DEFAULT') uses the correct module path and argument for loading pretrained weights.

3. Final Answer:

   model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights='DEFAULT') -> Option A

4. Quick Check:

   Use weights='DEFAULT' to load pretrained detection models [OK]

Hint: Use weights='DEFAULT' for pretrained detection models [OK]

Common Mistakes:

• Using pretrained=True instead of weights='DEFAULT'
• Wrong module path
• Missing 'models' in import path

3. What will be the output type of the following code snippet?

import torch
import torchvision
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights='DEFAULT')
model.eval()
inputs = [torch.rand(3, 300, 400)]
outputs = model(inputs)
print(type(outputs))

medium

A.

B.

C.

D.

Solution

1. Step 1: Understand model output format

   torchvision detection models return a list of dictionaries, one per input image.

2. Step 2: Check the type of outputs variable

   Since inputs is a list with one tensor, outputs will be a list containing one dictionary.

3. Final Answer:

   <class 'list'> -> Option D

4. Quick Check:

   Detection model output = list of dicts [OK]

Hint: Detection models output a list, one dict per input image [OK]

Common Mistakes:

• Expecting a single tensor output
• Confusing output with dict directly
• Not wrapping inputs in a list

4. Identify the error in this code snippet that tries to use a pretrained Mask R-CNN model:

import torchvision
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
model.eval()

medium

A. The function name should be fasterrcnn_resnet50_fpn

B. Mask R-CNN is not available in torchvision

C. The argument 'pretrained' is deprecated; use 'weights' instead

D. Missing import for torch

Solution

1. Step 1: Check function and argument usage

   Mask R-CNN is available as maskrcnn_resnet50_fpn in torchvision.models.detection.

2. Step 2: Identify deprecated argument

   Using pretrained=True is deprecated; the correct way is weights='DEFAULT'.

3. Final Answer:

   The argument 'pretrained' is deprecated; use 'weights' instead -> Option C

4. Quick Check:

   Use weights='DEFAULT' instead of pretrained=True [OK]

Hint: Replace pretrained=True with weights='DEFAULT' [OK]

Common Mistakes:

• Using deprecated pretrained argument
• Confusing model names
• Assuming Mask R-CNN is missing

5. You want to fine-tune a pretrained RetinaNet model from torchvision on your custom dataset with 5 classes. Which step is necessary before training?

hard

A. Replace the model's box predictor head to output 6 classes (5 + background)

B. Change the input image size to 224x224

C. Set pretrained=False to train from scratch

D. Freeze all layers except the backbone

Solution

1. Step 1: Understand class output adjustment

   Detection models include a background class, so for 5 classes, output head must predict 6 classes.

2. Step 2: Modify the box predictor head

   Replace the box predictor to match the number of classes (5 + 1 background) before training.

3. Final Answer:

   Replace the model's box predictor head to output 6 classes (5 + background) -> Option A

4. Quick Check:

   Adjust output head for custom classes [OK]

Hint: Add 1 for background class when changing output head [OK]

Common Mistakes:

• Not adjusting output head for new classes
• Changing input size unnecessarily
• Freezing all layers without fine-tuning head