python pytorch dataset dataloader and custom datasets: Pract

python pytorch dataset dataloader and custom datasets requires a clear understanding of the core syntax, runtime behavior, and practical implementation patterns demonstrated in the examples below.

Training a PyTorch model means moving data from disk into batches that the training loop can consume, and the torch.utils.data.Dataset and torch.utils.data.DataLoader classes define that pipeline. When your data does not match the standard formats bundled with torchvision.datasets, you write a custom Dataset that knows how to locate and load individual samples, and DataLoader handles batching, shuffling, and parallel loading. Working with python pytorch dataset dataloader and custom datasets comes down to understanding the contract Dataset must satisfy and how DataLoader consumes it.

The Dataset contract and what DataLoader adds

A Dataset is an indexable object. PyTorch requires exactly two methods:

__len__ returns the number of samples in the dataset.
__getitem__(index) returns the sample at that index.

A sample can be a single tensor, or a tuple of tensors such as (input, label). DataLoader accepts any Dataset and produces batches by calling __getitem__ repeatedly, grouping the results, and optionally shuffling them.

The separation of responsibilities matters. The Dataset knows where data lives and how to load one sample. The DataLoader knows how to iterate, batch, shuffle, and parallelize. You can reuse the same Dataset with different DataLoader configurations, and you can swap datasets without changing the training loop.

Implementing the three required methods

A custom dataset typically has three methods, even though only two are required.

__init__ stores the metadata needed to locate samples: file paths, labels, or a parsed index. This method should be cheap. Loading every sample into memory here defeats the purpose of a lazy pipeline and can exhaust RAM on large datasets.

__len__ returns the length of the index structure. It must match the number of valid indices that __getitem__ accepts. A mismatch produces an IndexError when DataLoader samples an index beyond the valid range.

__getitem__ does the real work. It reads the sample at the given index, applies any transforms, and returns tensors. Because DataLoader calls this method once per sample, heavy I/O happens here, which is exactly what allows parallel workers to overlap disk reads with GPU computation.

A concrete custom dataset example

A common case is tabular data stored in a CSV file. The dataset below reads the file once in __init__ to build an index, then converts rows to tensors lazily in __getitem__.

import torch
from torch.utils.data import Dataset
import pandas as pd

class CSVDataset(Dataset):
    def __init__(self, csv_path, transform=None):
        self.df = pd.read_csv(csv_path)
        self.transform = transform

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        features = torch.tensor(
            row.drop("label").to_numpy(), dtype=torch.float32
        )
        label = torch.tensor(row["label"], dtype=torch.long)
        if self.transform is not None:
            features = self.transform(features)
        return features, label

__getitem__ returns a tuple of two tensors. DataLoader collates these into a batch where the first element is a tensor of shape (batch_size, num_features) and the second is a tensor of shape (batch_size,). The transform argument is optional and applied per sample, which keeps preprocessing consistent whether you train, validate, or test with the same dataset.

Adding transforms and preprocessing

Transforms are callables that take a sample and return a modified sample. For image data, torchvision.transforms provides ToTensor, Normalize, Resize, and Compose to chain them.

from torchvision import transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5]),
])

Transforms must be applied inside __getitem__, not to the whole batch. Each worker process loads and transforms samples independently, so batch-level transforms would break the per-sample contract and complicate the collation step. For non-image data, a transform is just any callable; you can write a small function that standardizes features or augments inputs.

How DataLoader batches and parallelizes

DataLoader wraps the dataset and controls iteration.

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
)

batch_size controls how many samples are grouped per iteration. shuffle=True randomizes sample order each epoch, which is important for stochastic gradient descent. num_workers spawns separate processes, each holding its own copy of the dataset, so sample loading happens concurrently with training. pin_memory=True allocates page-locked memory, which speeds up the transfer of batches to a GPU.

With shuffle=True and multiple workers, PyTorch uses a sampler to assign indices to workers. Each worker prefetches samples in the background, so the training loop rarely blocks on disk I/O when the worker count is reasonable.

Common failure modes

Several errors appear repeatedly when people first build custom datasets.

Returning samples with inconsistent shapes breaks collation. If one sample is a 3-channel image and another has 1 channel, DataLoader cannot stack them into a batch tensor and raises an error during collation. Every sample must have the same shape and dtype.

An incorrect __len__ causes IndexError. If __len__ reports more samples than __getitem__ can handle, the sampler eventually requests an invalid index.

On Windows, multiprocessing with num_workers > 0 requires the dataset creation and training code to be inside an if __name__ == "__main__": guard, because worker processes re-import the main module. Without the guard, workers recursively spawn new processes.

Setting num_workers too high for a small dataset adds process-spawning overhead without benefit. The workers spend more time starting up than loading samples.

Forgetting to convert data to tensors is another common issue. If __getitem__ returns lists of varying length, DataLoader cannot collate them into a tensor. Convert to tensors with a fixed shape inside __getitem__.

Performance and memory considerations

The dominant cost in a data pipeline is usually I/O, not the DataLoader loop itself. num_workers is the primary lever: more workers overlap disk reads with GPU computation. The right value depends on the storage system and dataset size, and it is worth testing a few values rather than assuming more is better.

For small datasets that fit in memory, caching inside __getitem__ avoids repeated disk reads. A simple dictionary keyed by index stores loaded samples after the first access. This trades memory for I/O and is only sensible when the dataset is small enough to hold in RAM.

pin_memory=True is a low-cost improvement for GPU training. It allocates batches in pinned memory, which the CUDA runtime can copy to the device without an extra staging step. On CPU-only training it has no effect.

Newer PyTorch versions expose prefetch_factor, which controls how many batches each worker prefetches ahead of time. Increasing it can hide latency on slow storage, but it also increases memory usage per worker. The default is usually sufficient unless profiling shows the training loop waiting on data.

Choosing between built-in datasets and a custom implementation

torchvision.datasets provides ready-made classes for common formats: ImageFolder for directory-structured images, MNIST, CIFAR10, and others. If your data matches one of these layouts, using the built-in class saves code and is well tested.

Write a custom Dataset when the data lives in a proprietary format, requires domain-specific preprocessing, streams from a database, or is generated on the fly. The custom class gives you full control over how each sample is loaded and transformed, and it keeps that logic in one place rather than scattered through the training script.

The decision comes down to fit. If a built-in class matches your data layout, prefer it. If you need control over loading, caching, or augmentation, a custom Dataset is the right tool, and DataLoader works with it unchanged.