Pre-trained models (ResNet, VGG, EfficientNet) in Computer Vision

Introduction

Pre-trained models help us use smart image recognition without training from scratch. They save time and work well on many tasks.

You want to recognize objects in photos quickly without building a model yourself.

You have a small dataset and need a strong model that already learned from many images.

You want to try different models to see which works best for your pictures.

You need to build an app that identifies things like animals, cars, or faces fast.

You want to improve your model by starting from a good base instead of random guesses.

Syntax

Computer Vision

from torchvision import models

# Load a pre-trained model
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)

# Use the model for prediction or fine-tuning

Use weights=models.ResNet50_Weights.DEFAULT to load weights learned from large datasets like ImageNet.

You can replace resnet50 with vgg16 or efficientnet_b0 for other models.

Examples

Load ResNet50, a deep model good for many image tasks.

Computer Vision

from torchvision import models
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)

Load VGG16, a simpler but effective model for image recognition.

Computer Vision

from torchvision import models
model = models.vgg16(weights=models.VGG16_Weights.DEFAULT)

Load EfficientNet B0, a newer model that balances speed and accuracy.

Computer Vision

from torchvision import models
model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)

Sample Model

This code loads a pre-trained ResNet50 model and runs a dummy red image through it. It prints the predicted class index from ImageNet classes.

Computer Vision

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

# Load a pre-trained ResNet50 model
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.eval()  # Set to evaluation mode

# 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])
])

# Load an example image (replace 'image.jpg' with your image path)
img = Image.new('RGB', (224, 224), color='red')  # Create a red image for demo
input_tensor = preprocess(img)
input_batch = input_tensor.unsqueeze(0)  # Create batch dimension

# Run the model to get predictions
with torch.no_grad():
    output = model(input_batch)

# Get predicted class index
_, predicted_idx = torch.max(output, 1)

print(f"Predicted class index: {predicted_idx.item()}")

Important Notes

Pre-trained models are trained on ImageNet, which has 1000 classes of common objects.

You can fine-tune these models by training on your own data to improve accuracy for your task.

Make sure to preprocess images correctly to match the model's expected input.

Summary

Pre-trained models save time by using knowledge from large datasets.

ResNet, VGG, and EfficientNet are popular choices with different strengths.

Use them to quickly build image recognition apps or improve your own models.