PyTorch Pretrained Models

TorchVision ships with dozens of pretrained models ready to use in two lines of code. These models were trained on ImageNet — 1.2 million images across 1,000 categories. You can use them as-is for feature extraction, plug them into classification pipelines, or fine-tune them for custom tasks.

Loading a Pretrained Model

import torchvision.models as models

# New-style (TorchVision ≥ 0.13): use weights parameter
from torchvision.models import ResNet18_Weights

model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.eval()   # set to evaluation mode for inference

Making Predictions on a Single Image

import torch
from torchvision import transforms
from PIL import Image

# Use the preprocessing the model was trained with
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])
])

img = Image.open('cat.jpg').convert('RGB')
x   = preprocess(img).unsqueeze(0)   # shape: [1, 3, 224, 224]

with torch.no_grad():
    logits = model(x)
    probs  = torch.softmax(logits, dim=1)
    top5   = probs.topk(5)

print(top5.indices)   # top 5 class indices
print(top5.values)    # top 5 probabilities

Available Pretrained Models in TorchVision

Classification:
  ResNet family:     resnet18, resnet34, resnet50, resnet101, resnet152
  EfficientNet:      efficientnet_b0 through efficientnet_b7
  MobileNet:         mobilenet_v2, mobilenet_v3_small, mobilenet_v3_large
  VGG:               vgg11, vgg13, vgg16, vgg19
  Vision Transformer: vit_b_16, vit_b_32, vit_l_16
  ConvNeXt:          convnext_tiny, small, base, large

Detection:
  Faster R-CNN:      fasterrcnn_resnet50_fpn
  RetinaNet:         retinanet_resnet50_fpn
  SSD:               ssd300_vgg16

Segmentation:
  DeepLabV3:         deeplabv3_resnet50, deeplabv3_resnet101
  FCN:               fcn_resnet50, fcn_resnet101

Inspecting Model Architecture

model = models.resnet18(weights=None)   # no pretrained weights
print(model)

# Count parameters
total = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total:,}")   # 11,689,512 for ResNet18

Replacing the Classification Head

Every classification model ends with a fully connected layer sized for ImageNet's 1,000 classes. Replace it with a layer sized for your number of classes.

import torch.nn as nn
import torchvision.models as models

model = models.resnet50(weights='IMAGENET1K_V2')

# Identify the final layer
print(model.fc)   # Linear(in_features=2048, out_features=1000, bias=True)

# Replace with your task's layer
model.fc = nn.Linear(2048, 4)   # 4-class problem

# Verify
print(model.fc)   # Linear(in_features=2048, out_features=4, bias=True)

Different Architectures — Different Head Names

Architecture      Final layer attribute
──────────────────────────────────────────
ResNet            model.fc
EfficientNet      model.classifier[1]
MobileNetV3       model.classifier[3]
VGG               model.classifier[6]
ViT               model.heads.head

Using EfficientNet

from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
import torch.nn as nn

model = efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)

# Replace classifier head (1280 → your classes)
model.classifier[1] = nn.Linear(1280, 10)

x = torch.rand(4, 3, 224, 224)
print(model(x).shape)   # torch.Size([4, 10])

Feature Extraction Mode — Remove the Head

Sometimes you want a pretrained backbone to output features (embeddings) without classification. Remove the final layer and use the output as a fixed-size vector representation of the input image.

import torch.nn as nn
import torchvision.models as models

backbone = models.resnet18(weights='IMAGENET1K_V1')
backbone.fc = nn.Identity()   # replace fc with passthrough

backbone.eval()
with torch.no_grad():
    x = torch.rand(8, 3, 224, 224)
    features = backbone(x)
    print(features.shape)   # torch.Size([8, 512])
    # Each image is now represented as a 512-dimensional vector

Normalizing Input Correctly

Every ImageNet model expects input normalized with ImageNet's channel statistics. Using the wrong normalization is one of the most common transfer learning mistakes — the model behaves as if seeing garbage input.

ImageNet normalization:
  mean = [0.485, 0.456, 0.406]   ← per RGB channel
  std  = [0.229, 0.224, 0.225]

Always apply this transform after ToTensor():
  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std= [0.229, 0.224, 0.225])

Summary

TorchVision provides dozens of pretrained models for classification, detection, and segmentation. Load them with the weights parameter. Replace the final classification layer with one sized for your number of classes — the attribute name varies by architecture (model.fc for ResNet, model.classifier for EfficientNet). Always use the same input preprocessing (resize, crop, normalize) that the model was trained with. Use nn.Identity() to strip the classification head and extract feature vectors. Pretrained models dramatically reduce training time and improve accuracy on small datasets by starting from rich, already-learned representations.