torchvision Datasets and the v2 Transforms API

Built-in datasets and the modern transforms pipeline

torchvision ships ready-to-use datasets (MNIST, CIFAR-10/100, ImageNet, COCO) and a transforms pipeline. The current API is torchvision.transforms.v2 — it replaced the older torchvision.transforms as the recommended path. Use v2 for new code.

Why v2 matters

Native support for tv_tensors (Image, BoundingBoxes, Mask, Video) — important for detection and segmentation.
Transforms apply correctly across multiple inputs simultaneously: rotate an image AND its bounding boxes AND its mask, in one pass.
Much faster on the most common ops thanks to a rewritten implementation.
Cleaner separation between PIL-image transforms and tensor transforms.

The standard preprocessing chain

For ImageNet-pretrained models you'll see this exact recipe constantly:

ToImage() — wrap as a tv_tensor.Image (the v2-native type).
ToDtype(torch.float32, scale=True) — convert to float32 in [0, 1].
Resize(256) — resize the shorter side to 256.
CenterCrop(224) — crop the central 224x224.
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) — ImageNet normalization. Memorize these numbers.

Code

Built-in datasets — CIFAR10 example·python

import torch
import torchvision
from torchvision import datasets
import torchvision.transforms.v2 as T

transform = T.Compose([
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Resize(32),
    T.Normalize(mean=[0.4914, 0.4822, 0.4465],
                std=[0.2470, 0.2435, 0.2616]),  # CIFAR10-specific
])

train_ds = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
test_ds  = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
print(len(train_ds), len(test_ds))   # 50000 10000
print(train_ds[0][0].shape, train_ds[0][1])  # torch.Size([3, 32, 32]) 6

Standard ImageNet preprocessing·python

import torch
import torchvision.transforms.v2 as T

# This is the chain that matches every torchvision pretrained model
preprocess = T.Compose([
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Resize(256),
    T.CenterCrop(224),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]),
])

# But the modern recommended way is to ASK the model for its transforms:
from torchvision.models import resnet50, ResNet50_Weights
weights = ResNet50_Weights.IMAGENET1K_V2
preprocess = weights.transforms()
# weights.transforms() returns the EXACT preprocessing the model was trained with

v2's killer feature — multi-input transforms·python

import torch
import torchvision.transforms.v2 as T
from torchvision import tv_tensors

# A scene with image + bounding boxes + segmentation mask
img = torch.randint(0, 255, (3, 224, 224), dtype=torch.uint8)
boxes = tv_tensors.BoundingBoxes(
    [[10, 20, 100, 150]], format='XYXY', canvas_size=(224, 224)
)
mask = tv_tensors.Mask(torch.zeros(224, 224, dtype=torch.uint8))

transform = T.Compose([
    T.RandomHorizontalFlip(p=1.0),
    T.RandomRotation(15),
])

# Apply to all three at once — boxes and mask transform consistently with image
img_t, boxes_t, mask_t = transform(img, boxes, mask)
print(img_t.shape, boxes_t, mask_t.shape)
# This was nearly impossible with the old transforms API.

torchvision Datasets and the v2 Transforms API

Built-in datasets and the modern transforms pipeline

Why v2 matters

The standard preprocessing chain

Code

External links

Exercise

Progress

Comments 0