torchvision pre-trained models in PyTorch

Introduction

Pre-trained models help you use ready-made AI models that already learned from lots of images. This saves time and effort when building your own image tasks.

You want to quickly classify images without training a model from scratch.

You have a small dataset and need a model that already knows general image features.

You want to try out different AI models to see which works best for your images.

You want to use a model as a starting point and fine-tune it for your specific task.

You need a reliable model for tasks like object detection or segmentation using common architectures.

Syntax

PyTorch

import torchvision.models as models

model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.eval()

Set weights=models.ResNet18_Weights.DEFAULT to load a model with weights learned on ImageNet data.

Call model.eval() to set the model to evaluation mode before using it for predictions.

Examples

This loads ResNet18 with weights trained on ImageNet. Use it to classify images into 1000 classes.

PyTorch

import torchvision.models as models

# Load a pre-trained ResNet18 model
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.eval()

MobileNetV2 is a smaller, faster model good for mobile or limited hardware.

PyTorch

import torchvision.models as models

# Load a pre-trained MobileNetV2 model
model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT)
model.eval()

DenseNet121 is another popular model with dense connections for better feature reuse.

PyTorch

import torchvision.models as models

# Load a pre-trained DenseNet121 model
model = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
model.eval()

Sample Model

This program loads a pre-trained ResNet18 model, downloads an image of a dog, preprocesses it, and predicts the class with confidence.

PyTorch

import torch
from torchvision import models, transforms
from PIL import Image
import requests

# Load a pre-trained ResNet18 model
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.eval()

# Define image transforms to prepare input
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# Download an example image
url = 'https://upload.wikimedia.org/wikipedia/commons/9/9a/Pug_600.jpg'
image = Image.open(requests.get(url, stream=True).raw)

# Preprocess the image
input_tensor = preprocess(image)
input_batch = input_tensor.unsqueeze(0)  # create batch dimension

# Run the model on the input
with torch.no_grad():
    output = model(input_batch)

# Load ImageNet class names
labels_url = 'https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt'
labels = requests.get(labels_url).text.splitlines()

# Get the predicted class index
probabilities = torch.nn.functional.softmax(output[0], dim=0)
confidence, class_idx = torch.max(probabilities, dim=0)

# Print the result
print(f'Predicted class: {labels[class_idx]}')
print(f'Confidence: {confidence.item():.4f}')

Important Notes

Pre-trained models are trained on ImageNet with 1000 classes, so predictions match those classes.

Always preprocess images the same way the model expects (resize, crop, normalize).

Use model.eval() to turn off training features like dropout for correct predictions.

Summary

Pre-trained models let you use powerful image models without training from scratch.

They are great for quick experiments, small datasets, or as a starting point for your own tasks.

Remember to preprocess images correctly and set the model to evaluation mode before predicting.