What are Dataset and DataLoader in PyTorch and how do they work together?

PyTorch's data pipeline follows a clean two-class design: Dataset defines how to access a single sample (index → data), and DataLoader wraps a Dataset to handle batching, shuffling, and parallel loading.

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

class TabularDataset(Dataset):
    def __init__(self, X: np.ndarray, y: np.ndarray):
        # Convert once at construction — not inside __getitem__!
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

    def __len__(self) -> int:
        """Required — tells DataLoader how many samples exist."""
        return len(self.X)

    def __getitem__(self, idx: int):
        """Required — return a single (features, label) sample."""
        return self.X[idx], self.y[idx]

# Synthetic data
X = np.random.randn(1000, 20).astype(np.float32)
y = np.random.randint(0, 3, size=1000)

dataset = TabularDataset(X, y)
print(len(dataset))          # 1000
print(dataset[0])            # (tensor of 20 features, tensor scalar label)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,            # shuffle each epoch — essential for training
    num_workers=4,           # parallel data loading subprocesses
    pin_memory=True,         # faster CPU→GPU transfer
    drop_last=True,          # drop incomplete final batch
)

# Iterate over batches
for X_batch, y_batch in loader:
    print(X_batch.shape, y_batch.shape)  # (32, 20) (32,)
    break

# torchvision pre-built datasets
from torchvision import datasets, transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])
mnist = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
mnist_loader = DataLoader(mnist, batch_size=64, shuffle=True)

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.