Table of Contents
The Two Core Ideas: A Dataset and a DataLoader
In PyTorch, data handling usually starts by separating two responsibilities. A Dataset answers the question, what is the $i$th example, and how many examples exist. A DataLoader answers the question, how do we iterate over those examples efficiently in mini batches, possibly shuffled, possibly using multiple worker processes, and in a form that the training loop can consume.
This division is practical because you can swap datasets without changing your training code, and you can change batching or shuffling behavior without rewriting how samples are stored or decoded.
What a Dataset Is in PyTorch
A Dataset is a Python object that represents a collection of samples. Most commonly it behaves like an indexable container. Conceptually, it provides two things: a length and a way to fetch one sample by index. In code, this corresponds to implementing __len__ and __getitem__.
A single sample is typically returned as a tuple like (x, y) where x is the input tensor and y is the target, but it can also be a dictionary, or a richer structure if you need multiple inputs or metadata.
A Dataset should return one sample, not a batch. Batching is the DataLoader’s job.
Many real world datasets are not fully loaded into memory. A Dataset can read from disk on demand, decode images, parse rows from a CSV file, or look up items in an index. This is why __getitem__ is written as “fetch item $i$”, not “return everything”.
Map-Style Datasets vs Iterable Datasets
The most common type is a map style dataset, which supports random access by index. If your data can be indexed, you typically implement __len__ and __getitem__ and let the DataLoader handle shuffling and sampling.
Iterable datasets are designed for streaming data where random access is not available or not meaningful, such as reading from a large log file, a message queue, or a generator pipeline. In this case you implement __iter__ and yield samples one by one. Shuffling is not handled the same way for iterable datasets, so beginners generally start with map style datasets unless streaming is required.
If you want shuffle=True in a standard way, use a map style dataset. Iterable datasets require different shuffling strategies and are easier to get wrong.
Using Built-In TorchVision Datasets Quickly
For image tasks, torchvision.datasets provides ready made datasets such as MNIST and CIFAR. These typically download data and implement the Dataset interface for you. You only need to specify where to store data and whether you want train or test splits. The part that changes the most is the transform, which is covered in the next chapter, but it helps to understand that these datasets return a single sample per index and the DataLoader will assemble batches.
Writing a Custom Dataset
A custom dataset is appropriate when your data lives in your own folder structure, CSV files, or you have custom parsing logic. For a map style dataset, the minimal pattern is:
You store a list of references to samples, such as file paths, row indices, or already loaded arrays. __len__ returns how many references you have. __getitem__(idx) uses the reference at idx to load or construct x and y, converts them to tensors if needed, and returns them.
Two practical guidelines matter immediately for beginners. First, keep __getitem__ deterministic for a given index unless you are intentionally adding randomness through transformations. Second, keep heavy work inside __getitem__ reasonably fast, because it runs once per sample.
Do not call .to(device) inside your Dataset in most cases. Datasets should usually produce CPU tensors. Move batches to GPU in the training loop.
What a DataLoader Does
A DataLoader wraps a Dataset and provides an iterator that yields mini batches. Each step of iteration returns a batch of samples, usually as tensors stacked along a new first dimension, which is the batch dimension. If one sample is shaped like (C, H, W) for an image, then a batch of size $B$ is shaped like (B, C, H, W) after collation.
The DataLoader also handles shuffling, parallel workers, pinned memory, and how to combine samples into a batch.
The DataLoader yields batches. Your training loop should expect batch shaped tensors, not single examples.
Batch Size and the Shape Contract
The batch size controls how many samples are grouped per iteration. This affects memory usage, training speed, and optimization behavior, but those topics belong elsewhere. Here the key is the shape contract: your model and loss must accept batched inputs.
If your Dataset returns x and y as tensors, the DataLoader will typically stack x values and stack y values. This is why consistent shapes across samples matter.
When shapes are not consistent, such as variable length sequences, you need a custom collation strategy, which is handled in the dedicated chapter on collate functions. For now, assume that every sample in a batch has the same shape.
Shuffling and Sampling Basics
For training, you typically want shuffle=True so the model does not see samples in the same order every epoch. For validation and testing, you typically want shuffle=False so evaluation is stable and easier to compare.
Under the hood, shuffling is achieved by a sampler that produces indices in a shuffled order for map style datasets. If you later need class balancing or weighted sampling, that builds on the same idea, but it is covered in the chapter on sampling and class imbalance.
Use shuffle=True for training and shuffle=False for validation and test, unless you have a specific reason to do otherwise.
Workers, Prefetching, and Pin Memory
A DataLoader can use multiple subprocesses to load data in parallel using num_workers. This is important when loading data is slow compared to model computation. If num_workers=0, data loading happens in the main process.
Pinned memory, enabled with pin_memory=True, can speed up transferring CPU batches to CUDA GPUs because page locked memory allows faster DMA transfers. This does not help on pure CPU training, and behavior differs across platforms.
A good beginner default is to start with num_workers=0 while debugging, then increase it once the pipeline works. Increase gradually and watch for issues like deadlocks on some environments, or high RAM usage.
Start with num_workers=0 to debug correctness. Add workers later for speed, because multiprocessing can hide errors and make tracebacks harder to read.
The Default Collation Behavior
When the DataLoader builds a batch, it uses a function called collate_fn. If you do not provide one, PyTorch uses a default that stacks tensors, converts NumPy arrays to tensors, and groups nested structures in a predictable way.
This default works well when each sample has the same tensor shapes. When samples have different shapes, the default collation will fail. The solution is not to hack the Dataset to return padded batches, but to provide a custom collate_fn, which is covered in its own chapter.
A Minimal End-to-End Pattern
A common pattern is: create a Dataset, wrap it in a DataLoader for training, and wrap it in another DataLoader for validation. In your training loop, you iterate over the training DataLoader, get (x_batch, y_batch), move them to the chosen device, run the model, compute loss, and update parameters. The important point here is that the training loop should not care whether your Dataset reads images, CSV rows, or generated samples, as long as it provides samples consistently.
Common Beginner Mistakes with Datasets and DataLoaders
One frequent mistake is returning Python numbers or lists instead of tensors, which can lead to inconsistent dtypes or device issues later. Another is performing random shuffling inside the Dataset and also enabling shuffle=True, which makes the effective order harder to reason about and harms reproducibility.
A third mistake is putting expensive global work inside __getitem__, such as reading and parsing an entire CSV file for each sample, which can make training unbearably slow. Expensive one time setup should happen in __init__, while per sample work should be limited to what is needed for that sample.
Keep Dataset indexing simple: store references in __init__, load one sample in __getitem__, return tensors with consistent shapes. Let the DataLoader handle batching and shuffling.
Where This Fits in the Pipeline
At this point you should be able to recognize what belongs in a Dataset versus a DataLoader, choose map style for most beginner projects, and understand why your training loop expects batched tensors. Next, transforms and preprocessing extend this pipeline by modifying each sample, typically inside the Dataset via a transform callable, without changing the rest of your training code.