PyTorch Dataset Class

The Dataset class is the standard way to wrap your data in PyTorch. It teaches PyTorch how to access individual samples from your data — one item at a time. Once you have a Dataset, the DataLoader uses it to serve batches of samples to the training loop efficiently.

Why a Separate Dataset Class

Data comes in many forms: CSV files, image folders, database rows, audio files. PyTorch's Dataset class gives you a consistent interface regardless of where or how data is stored. Write the loading logic once in a Dataset, and the rest of the pipeline — DataLoader, batching, shuffling — works automatically.

Two Required Methods

Any Dataset subclass must implement:

__len__(self):
  → returns total number of samples in the dataset
  → DataLoader uses this to know when one epoch is done

__getitem__(self, idx):
  → receives an integer index
  → returns the sample at that position (usually a tuple of data + label)

Built-in Datasets

PyTorch and TorchVision provide many ready-to-use datasets. These are perfect for learning and benchmarking.

from torchvision import datasets, transforms

# Download MNIST (handwritten digits)
train_data = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

print(len(train_data))         # 60000 samples
print(train_data[0][0].shape)  # torch.Size([1, 28, 28]) — image
print(train_data[0][1])        # 5 — label (digit class)

Creating a Custom Dataset from a Tensor

The simplest custom dataset wraps two tensors — inputs and targets — into a Dataset.

import torch
from torch.utils.data import Dataset

class NumberDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __len__(self):
        return len(self.X)    # total number of samples

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]   # return one sample + its label

# Create dataset
X = torch.rand(200, 5)
y = torch.randint(0, 3, (200,))

dataset = NumberDataset(X, y)
print(len(dataset))           # 200
print(dataset[0][0].shape)    # torch.Size([5])
print(dataset[0][1])          # integer class label

TensorDataset – Shortcut for Simple Cases

When your data already lives in tensors, PyTorch offers TensorDataset — a built-in Dataset that wraps multiple tensors without you writing a class.

from torch.utils.data import TensorDataset

X = torch.rand(100, 4)
y = torch.rand(100, 1)

dataset = TensorDataset(X, y)
print(len(dataset))      # 100
print(dataset[0])        # (tensor([...]), tensor([...]))

Loading Images from a Folder

TorchVision's ImageFolder loads images organized in subfolders where each subfolder name is a class label.

data/
├── cats/
│   ├── cat001.jpg
│   └── cat002.jpg
└── dogs/
    ├── dog001.jpg
    └── dog002.jpg

from torchvision import datasets, transforms

dataset = datasets.ImageFolder(
    root='./data',
    transform=transforms.ToTensor()
)

print(dataset.classes)    # ['cats', 'dogs']
print(len(dataset))       # total image count

Splitting into Train and Validation Sets

from torch.utils.data import random_split

dataset = NumberDataset(X, y)

train_size = int(0.8 * len(dataset))   # 80% training
val_size   = len(dataset) - train_size  # 20% validation

train_set, val_set = random_split(dataset, [train_size, val_size])

print(f"Train: {len(train_set)}  Val: {len(val_set)}")

Dataset Flow Diagram

Raw Data Source (CSV / folder / array)
         │
         ▼
  Custom Dataset class
   __len__() → tells how many samples
   __getitem__(idx) → returns sample at idx
         │
         ▼
  DataLoader
   batches samples
   shuffles order
   runs prefetching on multiple CPU workers
         │
         ▼
  Training Loop (receives batch tensors)

Summary

The Dataset class wraps your data so PyTorch knows how to access it. Implement __len__ to return the sample count and __getitem__ to return one sample by index. For tensor data, use the built-in TensorDataset. For image folders, use ImageFolder from TorchVision. For CSV files, databases, or anything custom, subclass Dataset and write your own loading logic. Use random_split to divide a dataset into training and validation sets. The Dataset class is the entry point for all data in the PyTorch training pipeline.