Pretrained Models

A pretrained model is a neural network architecture distributed together with parameters learned on an earlier task. For computer vision models, the earlier task is often a large-scale image classification problem such as ImageNet. A pretrained model can be used directly for inference on new data when its original labels match our needs. More commonly, it can be used as a starting point for transfer learning, since the learned features are often generalizable for other tasks. Such portability is one of the key advantages of deep learning over traditional learning methods.

TorchVision provides pretrained models for classification, detection, segmentation, keypoints, and optical flow. Each set of weights is tied to a particular training recipe: label definitions, preprocessing, and network architecture. In this tutorial, we focus on the inference workflow of a pretrained model.

import torch
import torchvision
import matplotlib.pyplot as plt

Getting started

The torchvision.models API provides a simple way to instantiate predefined networks and download pretrained weights. The API allows loading different pretrained weights on the same model, keeps track of meta-data such as the classification labels, and includes the preprocessing transforms necessary to prepare the input data for the models. Keeping these pieces together prevents a common error, which is feeding an image to a model using preprocessing that does not match its training.

Cache directory

The first time a weight set is requested, TorchVision downloads it to the PyTorch Hub cache. Later calls reuse the cached file. By default, this directory is set to <hub_dir>/checkpoints, where <hub_dir> is the directory returned by torch.hub.get_dir(). Let’s check the cache directory on our system.

Show code cell source

Hide code cell source

print('Hub directory:', torch.hub.get_dir())

Hub directory: /Users/giovanni/.cache/torch/hub

There are two ways to change the Hub directory.

• Set the environment variable TORCH_HOME to the desired directory.

• Call the function torch.hub.set_dir() with the desired directory as argument.

We will keep the default cache directory for this tutorial.

Loading a pretrained model

For each available model, TorchVision provides a method to instantiate it with or without pretrained weights. The method signature is as follows.

torchvision.models.<model_name>(weights=<weights_option>)

Here, <model_name> is the name of a pretrained model, and <weights_option> specifies the weights to be loaded (or None for randomly initialized weights). As a general rule, the pretrained models are named in lower case (resnet50, mobilenet_v3_large, …), and the weights are named in camel case with the suffix _Weights (ResNet50_Weights, MobileNet_V3_Large_Weights, …). The names of available pretrained models are listed by the function torchvision.models.list_models().

For example, here is how to instantiate a ResNet-50 model with pretrained weights.

weights = torchvision.models.ResNet50_Weights.DEFAULT

model = torchvision.models.resnet50(weights=weights)

Using a pretrained model

Before we can use a pretrained model, we must prepare the input data in the same way as the training data. This includes resizing the image to the correct resolution, rescaling the pixel values, and applying any other preprocessing steps that were used during training. The specific preprocessing steps can vary between different models and even between different versions of the same model.

To simplify inference, TorchVision bundles the preprocessing transforms into each model weight. Calling transforms() on a weight object returns the inference preprocessing pipeline for that weight. Furthermore, some models use layers which have different training and evaluation behavior, such as batch normalization. To switch between these modes, we must use model.train() or model.eval() as appropriate.

model.eval()

preprocess = weights.transforms()

Note

Evaluation mode changes the behavior of layers such as dropout and batch normalization. But it does not disable gradient tracking by itself. During inference, we use torch.inference_mode() to avoid building an unnecessary computation graph and reduce memory use.

Classification

Image classification is a computer vision task that involves predicting the class of an image. TorchVision provides several pretrained models that can be used for image classification. These models are trained on the ImageNet-1K dataset, which consists of 1.2 million images from 1000 classes. The table of all available classification weights can be found here, along with their performance metrics.

Loading a pretrained model

Here is how to instantiate a pretrained model for image classification. We need to load the pretrained weights into the selected architecture, set the model to evaluation mode, and retrieve the preprocessing transforms from the weights.

from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights

weights = MobileNet_V3_Large_Weights.DEFAULT

classifier = mobilenet_v3_large(weights=weights)
classifier.eval()

preprocess = weights.transforms()

Preparing test data

To demonstrate the model’s prediction capabilities, we are going to download the test images of Imagenette dataset, which is a subset of 10 easily classified classes from Imagenet (tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, parachute). The images come in three different resolutions (full, 320px, 160px).

from torchvision.datasets import Imagenette

test_ds = Imagenette('.data/Imagenette', split='val', download=True, size="160px")

print("Number of images:", len(test_ds))
print("Image size:", test_ds[0][0].size)

100%|██████████| 99.0M/99.0M [00:03<00:00, 28.3MB/s]

Number of images: 3925
Image size: (213, 160)

Making predictions

Let’s grab a batch of images from the test dataset, preprocess them, and make predictions with the pretrained model. The classifier returns a logit vector for each image, where each logit corresponds to a class in the ImageNet dataset. Logits are unrestricted scores, so we apply softmax to obtain non-negative values that sum to one. These values are useful for ranking classes, but they should not automatically be interpreted as calibrated probabilities: a model can assign a high softmax score and still be wrong

torch.manual_seed(14)

# Random selection of 4 test images
n = 4
idx = torch.randint(0, len(test_ds), (n,))
images, labels = zip(*[test_ds[i] for i in idx])

# Preprocessing and batching
batch = torch.stack([preprocess(img) for img in images])

# Inference without gradient tracking
with torch.inference_mode():
    logits = classifier(batch)
    probabilities = torch.nn.functional.softmax(logits, dim=1)

Let’s visualize the predictions generated by the model. We display the five classes with the largest softmax values. Top-k output is often more informative than showing only the top class, because it reveals plausible alternatives. The names in weights.meta["categories"] contains the class labels of ImageNet.

Show code cell source

Hide code cell source

categories = weights.meta["categories"]

for i, (img, label) in enumerate(zip(images, labels)):
    plt.figure(figsize=(6, 3), tight_layout=True)
    plt.subplot(1, 2, 1)
    plt.imshow(img)
    plt.axis('off')
    plt.title(f"{test_ds.classes[label][0]}")

    values, indices = torch.topk(probabilities[i], k=5)
    top5_cats = [categories[index] for index in indices.tolist()]

    plt.subplot(1, 2, 2)
    plt.barh(top5_cats, values.cpu())
    plt.gca().invert_yaxis()
    plt.xlabel('Softmax value')
    plt.title('Top-5 predictions')
    plt.show()

Object detection

Object detection is a computer vision task that involves locating and classifying objects in an image. TorchVision provides several pretrained detectors trained on COCO dataset, whose standard label set contains 80 object categories. A pretrained detector can only name categories represented by its training label space. An unfamiliar object may be missed or confused with a known class.

Loading a pretrained model

Here is how to instantiate a pretrained model for object detection. We need to load the pretrained weights into the selected architecture, set the model to evaluation mode, and retrieve the preprocessing transforms from the weights.

from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_fpn, FasterRCNN_MobileNet_V3_Large_FPN_Weights

weights = FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT

detector = fasterrcnn_mobilenet_v3_large_fpn(weights=weights)
detector.eval()

preprocess = weights.transforms()

Downloading: "https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth" to /Users/giovanni/.cache/torch/hub/checkpoints/fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth

100%|██████████| 74.2M/74.2M [00:01<00:00, 75.5MB/s]

Preparing example data

To demonstrate the model’s prediction capabilities, we are going to use images from the Pascal VOC dataset. VOC and COCO do not have identical label sets, so this is a demonstration of inference rather than a valid benchmark of the COCO-trained detector. Some VOC objects have corresponding COCO categories and others do not.

from torchvision.datasets import VOCDetection

test_ds = VOCDetection('.data/voc', year='2007', image_set='test', download=True)

print("Number of images:", len(test_ds))
print("Image size:", test_ds[0][0].size)

100%|██████████| 451M/451M [01:10<00:00, 6.42MB/s] 

Number of images: 4952
Image size: (353, 500)

Making predictions

In evaluation mode, the detector expects a list of image tensors in channel-first format, and returns one dictionary per image with the following keys.

• boxes - The bounding box coordinates of each detected object.

• labels - The class label of each detected object.

• scores - The ranking score of each detected object.

The detector already applies its model-specific postprocessing, including non-maximum suppression (NMS), which removes many highly overlapping duplicate boxes. We still choose a score threshold for display. That threshold controls a precision–recall tradeoff and should be selected on validation data for a real application; it is not a universal definition of confidence.

torch.manual_seed(10)

# # Random selection of 4 test images
n = 4
idx = torch.randint(0, len(test_ds), (n,))
images, labels = zip(*[test_ds[i] for i in idx])

# Preprocessing and batching
batch = [preprocess(img) for img in images]

# Inference without gradient tracking
with torch.inference_mode():
    predictions = detector(batch)

Let’s visualize the predictions generated by the model. For each image, we only display the boxes with a confidence score greater than 0.5. This value is only an illustrative display threshold; changing it affects which boxes are shown.

Show code cell source

Hide code cell source

from torchvision.utils import draw_bounding_boxes
from torchvision.transforms.v2.functional import to_image

threshold = 0.5

for pil_image, prediction in zip(images, predictions):
    accepted = prediction["scores"] > threshold
    label_ids = prediction["labels"][accepted]
    class_names = [weights.meta["categories"][i] for i in label_ids.tolist()]
    boxes = prediction["boxes"][accepted]

    tensor_image = to_image(pil_image)
    image_boxes = draw_bounding_boxes(
        tensor_image,
        boxes=boxes,
        labels=class_names,
        width=2,
        font="Arial",
        font_size=25,
        colors="red",
    )

    plt.figure(figsize=(4, 4))
    plt.imshow(image_boxes.permute(1, 2, 0))
    plt.axis('off')
    plt.show()

Summary

In this tutorial, we learned how to download TorchVision pretrained models and make predictions for image classification and object detection tasks. We also discussed how to use the preprocessing transforms provided by the weights to prepare input data for the models.