torchvision

The companion library to PyTorch for computer vision. torchvision provides image transforms (augmentation + normalization), pretrained models (ResNet, EfficientNet, ViT, Faster R-CNN), and ready-made dataset loaders — everything you need to build and fine-tune vision models in PyTorch.

Table of Contents

• Installation / Setup
• Core Syntax
• Common Operations
• Practical Examples
• Common Patterns
• Best Practices
• Common Mistakes
• Troubleshooting
• Related Cheat Sheets
• Official Documentation

Installation / Setup

pip install torchvision                    # or: uv add torchvision

Key imports:

import torchvision
from torchvision import transforms, datasets, models
from torchvision.transforms import v2
import torch

Core Syntax

Transforms — the preprocessing pipeline

from torchvision import transforms

# training: augment + normalize
train_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),                         # PIL → Tensor, [0,1], CHW
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),  # ImageNet stats
])

# validation: no augmentation
val_transform = 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]),
])

Transform	Purpose
Resize(size)	Scale shortest edge
CenterCrop(size)	Crop center square
RandomCrop(size)	Random crop (augmentation)
RandomHorizontalFlip(p=0.5)	Left-right flip
RandomVerticalFlip(p=0.5)	Top-bottom flip
RandomRotation(degrees)	Random rotation
ColorJitter(b, c, s, h)	Random brightness/contrast/saturation/hue
GaussianBlur(kernel)	Random blur
RandomAffine(degrees, translate)	Random affine
RandomResizedCrop(size, scale)	Crop + resize (efficient augmentation)
ToTensor()	PIL → FloatTensor [0,1], CHW
Normalize(mean, std)	Channel-wise standardization

v2 transforms (modern API)

from torchvision.transforms import v2

transform = v2.Compose([
    v2.Resize(256),
    v2.RandomCrop(224),
    v2.RandomHorizontalFlip(p=0.5),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

v2 transforms handle both images and bounding boxes — useful for detection.

Pretrained models

from torchvision import models

# image classification (pretrained on ImageNet)
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)
model = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT)

# object detection
model = models.detection.fasterrcnn_resnet50_fpn(weights=models.FasterRCNN_ResNet50_FPN_Weights.DEFAULT)

# set to eval mode for inference
model.eval()

Model	Type	Params	ImageNet top-1
resnet18	Classifier	11M	69.8%
resnet50	Classifier	25M	76.1%
efficientnet_b0	Classifier	5M	77.1%
vit_b_16	Classifier	86M	81.1%
fasterrcnn_resnet50_fpn	Detector	41M	—

Datasets

from torchvision import datasets
from torch.utils.data import DataLoader

train_ds = datasets.ImageFolder("data/train", transform=train_transform)
val_ds   = datasets.ImageFolder("data/val",   transform=val_transform)

print(train_ds.class_to_idx)     # {'cat': 0, 'dog': 1}
print(train_ds[0][0].shape)      # torch.Size([3, 224, 224])

train_dl = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)

ImageFolder expects a directory structure:

data/
  train/
    cat/  → img1.jpg, img2.jpg, ...
    dog/  → img1.jpg, img2.jpg, ...
  val/
    cat/  → ...
    dog/  → ...

Common Operations

ImageNet inference

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

model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT).eval()

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

img = transform(Image.open("dog.jpg").convert("RGB")).unsqueeze(0)   # add batch dim
with torch.no_grad():
    logits = model(img)
    probs = torch.nn.functional.softmax(logits, dim=1)
    top5 = torch.topk(probs, 5)

# decode labels
weights = models.ResNet50_Weights.DEFAULT
labels = weights.meta["categories"]
for i, (prob, idx) in enumerate(zip(top5.values[0], top5.indices[0])):
    print(f"{labels[idx]:30s} {prob:.3f}")

Fine-tune a pretrained model

import torch
from torchvision import models

model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)

# freeze early layers
for param in model.parameters():
    param.requires_grad = False

# replace final layer for your number of classes
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

# only train the new layer
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

Transfer learning with full fine-tuning

model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

# lower learning rate for pretrained layers, higher for new layer
optimizer = torch.optim.Adam([
    {"params": model.layer4.parameters(), "lr": 1e-5},
    {"params": model.fc.parameters(), "lr": 1e-3},
])

Denormalize (for visualization)

mean = torch.tensor([0.485, 0.456, 0.406])
std  = torch.tensor([0.229, 0.224, 0.225])

def denormalize(tensor):
    return tensor * std[:, None, None] + mean[:, None, None]

# convert to displayable numpy
img = denormalize(tensor.squeeze()).permute(1, 2, 0).clamp(0, 1).numpy()

Practical Examples

Example 1 — Train an image classifier

import torch
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# data
train_dl = DataLoader(
    datasets.ImageFolder("data/train", transform=transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])), batch_size=32, shuffle=True, num_workers=4)
val_dl = DataLoader(
    datasets.ImageFolder("data/val", transform=transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])), batch_size=32, num_workers=4)

# model
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.fc = torch.nn.Linear(model.fc.in_features, len(train_dl.dataset.classes))
model = model.to(device)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(5):
    model.train()
    total, correct = 0, 0
    for imgs, labels in train_dl:
        imgs, labels = imgs.to(device), labels.to(device)
        outputs = model(imgs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        correct += (outputs.argmax(1) == labels).sum().item()
        total += labels.size(0)
    print(f"Epoch {epoch+1}: train acc {correct/total:.3f}")

Example 2 — Faster R-CNN object detection

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

model = models.detection.fasterrcnn_resnet50_fpn(
    weights=models.FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
model.eval()

transform = transforms.Compose([
    transforms.ToTensor(),
])

img = transform(Image.open("street.jpg").convert("RGB")).unsqueeze(0)
with torch.no_grad():
    preds = model(img)[0]

# filter confident predictions
mask = preds["scores"] > 0.5
boxes = preds["boxes"][mask].numpy()
labels = preds["labels"][mask].numpy()
scores = preds["scores"][mask].numpy()

COCO_LABELS = {1: "person", 2: "bicycle", 3: "car", 4: "motorcycle", 8: "truck"}
for box, label, score in zip(boxes, labels, scores):
    name = COCO_LABELS.get(label, f"cls_{label}")
    print(f"{name}: {score:.2f} at {box.astype(int)}")

Common Patterns

Pattern: quick inference on a single image

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

model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT).eval()
weights = models.ResNet50_Weights.DEFAULT

img = Image.open("cat.jpg")
pred = weights.transforms()(img).unsqueeze(0)
with torch.no_grad():
    label = model(pred).argmax(1).item()
print(weights.meta["categories"][label])

Pattern: freeze + unfreeze schedule

# first train only the head
for param in model.parameters():
    param.requires_grad = False
model.fc.requires_grad_(True)

# after a few epochs, unfreeze everything
for param in model.parameters():
    param.requires_grad = True

Pattern: custom dataset with transforms

from torch.utils.data import Dataset
from PIL import Image

class MyDataset(Dataset):
    def __init__(self, paths, labels, transform=None):
        self.paths = paths
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, self.labels[idx]

Best Practices

• Always normalize with the same mean/std the pretrained model expects (ImageNet stats for all standard models).
• Use model.eval() and torch.no_grad() for inference.
• transforms.ToTensor() converts PIL to float [0,1] CHW tensor; Normalize then standardizes per-channel.
• Start with a small learning rate (1e-4 to 1e-3) for fine-tuning; higher for new layers.
• Freeze pretrained layers first; unfreeze later for full fine-tuning.
• Use RandomResizedCrop(224, scale=(0.8, 1.0)) as the single augmentation that handles resize + crop + aspect ratio.
• num_workers=4 in DataLoader for parallel data loading.

Common Mistakes

Mistake	Fix
Forgetting model.eval() during inference	Batch norm and dropout behave differently in eval mode
Forgetting torch.no_grad()	Inference with gradients wastes memory and is slower
Wrong input shape (HWC instead of CHW)	transforms.ToTensor() handles this; manual tensors need .permute(2,0,1)
Using training transforms for validation	Use Resize + CenterCrop, not random augmentation
Wrong normalization stats	Match the pretrained model's expected mean/std
CUDA out of memory	Lower batch size, use mixed precision, or smaller model

Troubleshooting

Symptom	Cause / Fix
PIL.UnidentifiedImageError	Corrupt images in dataset — filter them out
Poor accuracy after fine-tuning	Learning rate too high/low; try unfreezing more layers
CUDA out of memory	Lower batch_size; use torch.cuda.empty_cache()
Training accuracy high but val low	Overfitting — add more augmentation, regularization, or data
Different accuracy from expected	Check normalization; ensure model.eval() is set

Related Cheat Sheets

• opencv.md — video, contours, webcam (not in torchvision)
• pillow.md — PIL for image I/O
• yolo — YOLO detection with Ultralytics
• ultralytics — Ultralytics API
• PyTorch / Deep Learning — PyTorch training loop
• NumPy — array ops
• matplotlib — visualizing images and results

Official Documentation

• torchvision docs — https://pytorch.org/vision/
• transforms reference — https://pytorch.org/vision/stable/transforms.html
• pretrained models — https://pytorch.org/vision/stable/models.html

Last Updated

2026-09-06 — Initial version.