How to Use torchvision Datasets in PyTorch

Use torchvision.datasets by importing the dataset class you want, specifying parameters like root for data location, train to choose train/test split, and transform for preprocessing. Then wrap it with torch.utils.data.DataLoader to load data in batches for training or evaluation.

Syntax

The basic syntax to use a torchvision dataset is:

• torchvision.datasets.DatasetName(root, train=True/False, transform=None, download=True/False)
• root: folder path to store or load the dataset.
• train: whether to load training data (True) or test data (False).
• transform: preprocessing steps like converting images to tensors or normalization.
• download: if True, downloads the dataset if not found locally.

After creating the dataset object, use torch.utils.data.DataLoader(dataset, batch_size, shuffle) to load data in batches.

python

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

# Define a transform to convert images to tensors
transform = transforms.ToTensor()

# Load the training set of MNIST dataset
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)

# Create a data loader to load data in batches
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

Example

This example loads the MNIST dataset, applies a transform to convert images to tensors, and iterates through one batch to show the shape of images and labels.

python

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

# Transform to convert images to tensor
transform = transforms.ToTensor()

# Load MNIST training dataset
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)

# DataLoader for batch loading
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# Get one batch of data
images, labels = next(iter(train_loader))

print(f'Batch image tensor shape: {images.shape}')
print(f'Batch labels tensor shape: {labels.shape}')

Output

Batch image tensor shape: torch.Size([32, 1, 28, 28])
Batch labels tensor shape: torch.Size([32])

Common Pitfalls

• Forgetting to set download=True when running the dataset for the first time causes errors if data is not present locally.
• Not applying transform to convert images to tensors will cause errors when feeding data to models.
• Mixing train=True and train=False datasets without clear separation can confuse training and evaluation phases.
• Not using DataLoader for batching can lead to inefficient training.

python

from torchvision import datasets

# Wrong: No download, no transform
try:
    dataset = datasets.MNIST(root='./data', train=True)
except Exception as e:
    print(f'Error: {e}')

# Right: With download and transform
from torchvision import transforms
transform = transforms.ToTensor()
dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)

Output

Error: Dataset not found or corrupted. You can use download=True to download it

Quick Reference

Key points to remember when using torchvision datasets:

• Always specify root folder for data storage.
• Use train=True for training data, train=False for test data.
• Apply transform to convert images to tensors and normalize.
• Set download=True to get data automatically if missing.
• Wrap dataset with DataLoader for batch processing and shuffling.

Key Takeaways

Use torchvision.datasets with proper root, train, transform, and download parameters.

Always convert images to tensors using transforms before feeding to models.

Use DataLoader to load data in batches and shuffle for training.

Set download=True to automatically get datasets if not present locally.

Keep training and test datasets separate by using train=True or False.