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.

Practice

1. What is the main advantage of using torchvision pre-trained models?

easy

A. They automatically improve your dataset quality.

B. They generate new images from text descriptions.

C. They reduce the size of your images.

D. They allow you to use powerful image models without training from scratch.

Solution

1. Step 1: Understand what pre-trained models do

   Pre-trained models are already trained on large datasets, so you don't need to train them from zero.

2. Step 2: Identify the main benefit

   This saves time and resources, letting you use powerful models quickly.

3. Final Answer:

   They allow you to use powerful image models without training from scratch. -> Option D

4. Quick Check:

   Pre-trained models = reuse trained weights [OK]

Hint: Pre-trained means ready to use without full training [OK]

Common Mistakes:

• Thinking pre-trained models improve data quality
• Confusing pre-trained models with image resizing
• Believing they generate images from text

2. Which of the following is the correct way to load a pre-trained ResNet18 model from torchvision?

easy

A. model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)

B. model = torchvision.resnet18(pretrained=True)

C. model = torchvision.models.resnet18(pretrained=False)

D. model = torchvision.models.load_resnet18(pretrained=True)

Solution

1. Step 1: Recall the updated torchvision syntax

   Since torchvision 0.13+, pre-trained weights are loaded using the 'weights' argument with a weights enum, not 'pretrained=True'.

2. Step 2: Identify the correct syntax for ResNet18

   Use torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1).

3. Final Answer:

   model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1) -> Option A

4. Quick Check:

   Use weights=enum, not pretrained=True [OK]

Hint: Use weights= argument, not pretrained=True [OK]

Common Mistakes:

• Using pretrained=False which doesn't load pre-trained weights
• Calling torchvision.resnet18 directly
• Using a non-existent load_resnet18 function

3. What will be the output shape of the following code snippet using a pre-trained ResNet18 model on a batch of 8 RGB images of size 224x224?

import torch
import torchvision.models as models
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.eval()
inputs = torch.randn(8, 3, 224, 224)
outputs = model(inputs)
print(outputs.shape)

medium

A. torch.Size([8, 3, 224, 224])

B. torch.Size([8, 1000])

C. torch.Size([8, 512])

D. torch.Size([1, 1000])

Solution

1. Step 1: Understand ResNet18 output size

   ResNet18 pre-trained on ImageNet outputs logits for 1000 classes, so output shape is (batch_size, 1000).

2. Step 2: Check input batch size and output shape

   Input batch size is 8, so output shape is (8, 1000).

3. Final Answer:

   torch.Size([8, 1000]) -> Option B

4. Quick Check:

   Batch size 8, 1000 classes output [OK]

Hint: Output shape = (batch_size, number_of_classes) [OK]

Common Mistakes:

• Confusing output with input image shape
• Expecting feature vector size instead of class logits
• Assuming batch size 1 output

4. You wrote this code to use a pre-trained model for prediction but get wrong results:

model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
inputs = torch.randn(1, 3, 224, 224)
outputs = model(inputs)

What is the likely mistake?

medium

A. Input tensor shape should be (3, 224, 224) without batch dimension.

B. You need to set weights=None to use pre-trained weights.

C. You forgot to call model.eval() before prediction.

D. You must convert inputs to numpy arrays before passing to model.

Solution

1. Step 1: Check model mode for prediction

   Pre-trained models must be set to evaluation mode with model.eval() to disable dropout and batch norm updates.

2. Step 2: Identify the missing step

   The code misses model.eval(), so outputs may be incorrect or inconsistent.

3. Final Answer:

   You forgot to call model.eval() before prediction. -> Option C

4. Quick Check:

   Set model.eval() before inference [OK]

Hint: Always call model.eval() before predicting [OK]

Common Mistakes:

• Not calling model.eval() before inference
• Wrong input tensor shape without batch
• Trying to convert tensors to numpy before model

5. You want to fine-tune a pre-trained ResNet18 model on your own 5-class dataset. Which of the following code snippets correctly replaces the final layer for this task?

hard

A. model.fc = torch.nn.Linear(in_features=512, out_features=5)

B. model.classifier = torch.nn.Linear(in_features=1000, out_features=5)

C. model.fc = torch.nn.Linear(in_features=2048, out_features=5)

D. model.output = torch.nn.Linear(in_features=512, out_features=1000)

Solution

1. Step 1: Identify the final layer of ResNet18

   ResNet18's final fully connected layer is model.fc with input features 512 and output 1000 classes.

2. Step 2: Replace final layer for 5 classes

   To fine-tune, replace model.fc with a new Linear layer with 512 inputs and 5 outputs.

3. Final Answer:

   model.fc = torch.nn.Linear(in_features=512, out_features=5) -> Option A

4. Quick Check:

   Replace model.fc with correct output size [OK]

Hint: Replace model.fc with Linear(512, number_of_classes) [OK]

Common Mistakes:

• Replacing wrong attribute like model.classifier
• Using wrong input feature size (2048 instead of 512)
• Not changing output features to dataset classes