torchvision detection models in PyTorch

Introduction

Detection models help computers find and label objects in pictures, like spotting cars or people.

You want to find all the faces in a photo.

You need to detect cars and pedestrians in street images for a self-driving car.

You want to count how many animals are in a wildlife photo.

You want to highlight objects in a video for security cameras.

Syntax

PyTorch

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

Use pretrained=True to load a model already trained on common objects.

Call model.eval() to set the model for evaluation (not training).

Examples

Loads a Faster R-CNN model with a ResNet-50 backbone, ready to detect objects.

PyTorch

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

Loads a Mask R-CNN model that detects objects and also creates masks showing object shapes.

PyTorch

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

Loads a RetinaNet model, which is fast and good for detecting many objects.

PyTorch

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

Sample Model

This program loads an image, applies a pretrained Faster R-CNN model to detect objects, and prints labels and confidence scores for detections above 50%.

PyTorch

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

# Load a sample image
img = Image.open("sample.jpg")

# Transform image to tensor
transform = T.Compose([T.ToTensor()])
img_tensor = transform(img)

# Load pretrained Faster R-CNN model
model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

# Run detection
with torch.no_grad():
    predictions = model([img_tensor])

# Print detected labels and scores
labels = predictions[0]['labels']
scores = predictions[0]['scores']
print("Detected objects and scores:")
for label, score in zip(labels, scores):
    if score > 0.5:
        print(f"Label: {label.item()}, Score: {score.item():.2f}")

Important Notes

Labels are numbers that correspond to object categories (like 1 = person, 3 = car).

You can find label names in COCO dataset documentation to understand what each number means.

Make sure your input images are converted to tensors and normalized if needed.

Summary

torchvision provides ready-to-use detection models like Faster R-CNN, Mask R-CNN, and RetinaNet.

These models help find and label objects in images with good accuracy.

Use pretrained models for quick results without training from scratch.

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