In PyTorch, what must a custom map-style Dataset implement, and what does DataLoader add?

juniormust knowfreq 80%

answer

two methods, then a wrapper
one sample, not a batch
length plus indexed access
shuffling and batching live outside
__len__ and __getitem__(idx)

basics

~10 s

A map-style Dataset implements len, returning the number of samples, and getitem(index), returning one sample (often a tensor/label tuple). DataLoader wraps it and produces shuffled, collated mini-batches, optionally loaded by worker processes.

solid answer

~40 s

Subclassing torch.utils.data.Dataset means implementing two methods: __len__, which reports how many samples exist, and __getitem__(idx), which returns one sample — typically a (features, label) tuple, but any structure the collate function can handle. All the per-sample work (reading a file, decoding an image, applying transforms) belongs in __getitem__, because that is the unit the DataLoader parallelizes.

DataLoader supplies everything around that: the iteration order (shuffle=True installs a RandomSampler, otherwise a SequentialSampler), grouping indices into batches of batch_size, calling collate_fn to stack the samples into batched tensors, drop_last to discard a short final batch, worker processes via num_workers, and pin_memory for faster host-to-device copies. The Dataset knows nothing about batching or shuffling — that separation is the whole design.

Subclassing `torch.utils.data.Dataset` means implementing two methods: `__len__`, which reports how many samples exist, and `__getitem__(idx)`, which returns **one** sample — typically a `(features, label)` tuple, but any structure the collate function can handle. All the per-sample work (reading a file, decoding an image, applying transforms) belongs in `__getitem__`, because that is the unit the DataLoader parallelizes. `DataLoader` supplies everything around that: the iteration order (`shuffle=True` installs a `RandomSampler`, otherwise a `SequentialSampler`), grouping indices into batches of `batch_size`, calling `collate_fn` to stack the samples into batched tensors, `drop_last` to discard a short final batch, worker processes via `num_workers`, and `pin_memory` for faster host-to-device copies. The Dataset knows nothing about batching or shuffling — that separation is the whole design.

code

python · 19 lines

python

import torch
from torch.utils.data import Dataset, DataLoader


class SquaresDataset(Dataset):
    def __init__(self, n):
        self.n = n

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        x = torch.tensor([float(idx)])
        return x, x**2


loader = DataLoader(SquaresDataset(100), batch_size=8, shuffle=True, drop_last=True)
xb, yb = next(iter(loader))
print(xb.shape, yb.shape)  # torch.Size([8, 1]) torch.Size([8, 1])

go deeper

for a junior

Be able to write the three-method skeleton from memory — __init__, __len__, __getitem__ — and state plainly that __getitem__ returns one sample while DataLoader does the batching and shuffling.

for a middle

Explain the division of labour: sampler chooses indices, dataset fetches one sample each, collate stacks them. Name shuffle, batch_size, drop_last and what happens when a sampler is supplied alongside shuffle.

for a senior

Show you know where the cost sits: cheap index in __init__, expensive I/O and decode in __getitem__ so workers can parallelize it, and lazily-opened file handles so the dataset survives multiprocessing.

for a principal

Own the interface decision — whether datasets return dicts or tuples, where transforms are configured, and how splits are constructed — because a shared parent dataset with one transform pipeline is how augmentation quietly leaks into validation across a whole team's code.

Two dataset styles

torch.utils.data offers two dataset flavours. A map-style dataset models a mapping from an index to a sample: you subclass torch.utils.data.Dataset and implement __getitem__ and __len__. An iterable-style dataset subclasses torch.utils.data.IterableDataset and implements __iter__ instead, modelling a stream with no random access. Nearly all image, tabular and small-corpus pipelines use the map-style form, because random access is what makes shuffling and index-based splitting trivial.

The map-style contract

__len__(self) returns an integer: how many samples the dataset holds. It must be cheap — samplers call it every time an iterator is created.

__getitem__(self, idx) receives one integer index and returns one sample. It is not required to return a tensor; it may return a tuple (image_tensor, label_int), a dict {"input_ids": ..., "label": ...}, or a nested structure. The default collate function knows how to recurse through tuples, lists, dicts and namedtuples, so whatever shape you pick is preserved with an extra leading batch dimension.

The base Dataset class supplies almost nothing: it defines __add__, which returns a ConcatDataset, and leaves __getitem__ for you. There is no hidden machinery — the value is in the convention that DataLoader, Subset, random_split and ConcatDataset all rely on.

Put the expensive work inside __getitem__, not __init__. __init__ should build a lightweight index — a list of file paths, an offsets table, a dataframe of labels — while __getitem__ does the read, the decode and the transform. That is what gets distributed across worker processes; anything eagerly loaded in __init__ is loaded once, in the parent, and then duplicated into every worker.

What DataLoader adds

Order. shuffle=True wraps the dataset in a RandomSampler; shuffle=False gives a SequentialSampler. You can pass your own sampler (for example WeightedRandomSampler for class balancing), but sampler and shuffle are mutually exclusive and passing both raises ValueError. batch_sampler goes one level further and yields whole lists of indices, which is how length-bucketing is implemented.
Batching. batch_size indices are drawn, __getitem__ is called for each, and the resulting samples are handed to collate_fn.
Collation. The default collate stacks same-shaped tensors along a new dimension 0, converts NumPy arrays and Python numbers to tensors, and recurses through containers.
Parallelism. num_workers > 0 moves fetching and collation into separate processes.
Transfer. pin_memory=True copies each batch into page-locked memory so the host-to-device copy can be asynchronous.
Edges. drop_last=True discards a final short batch; timeout bounds how long the main process waits on a worker.

A DataLoader is not an iterator, it is an iterable: each for batch in loader creates a fresh iterator, which is why every epoch reshuffles.

Frequent mistakes

Returning a batch from __getitem__. If your __getitem__ returns a tensor of shape (32, 3, 224, 224) and you also set batch_size=32, you get (32, 32, 3, 224, 224). One index, one sample.

Forgetting __len__. Shuffling needs the length; without it len(dataset) raises TypeError and the default sampler cannot be built.

dtype drift. The default collate converts a NumPy array with torch.as_tensor, preserving dtype. A np.float64 array becomes a torch.float64 tensor, and feeding that to a float32 nn.Linear raises a scalar-type error. Cast in __getitem__.

Non-picklable handles. An open HDF5 file, an LMDB environment or a database connection created in __init__ breaks or corrupts under multiprocessing. Open them lazily, on first use inside the worker.

Shared transforms across splits. random_split returns Subset views over the same parent dataset, so the parent's augmentation pipeline applies to the validation subset too. Build two dataset objects with different transforms instead.

## Two dataset styles `torch.utils.data` offers two dataset flavours. A **map-style** dataset models a mapping from an index to a sample: you subclass `torch.utils.data.Dataset` and implement `__getitem__` and `__len__`. An **iterable-style** dataset subclasses `torch.utils.data.IterableDataset` and implements `__iter__` instead, modelling a stream with no random access. Nearly all image, tabular and small-corpus pipelines use the map-style form, because random access is what makes shuffling and index-based splitting trivial. ## The map-style contract `__len__(self)` returns an integer: how many samples the dataset holds. It must be cheap — samplers call it every time an iterator is created. `__getitem__(self, idx)` receives one integer index and returns **one sample**. It is not required to return a tensor; it may return a tuple `(image_tensor, label_int)`, a dict `{"input_ids": ..., "label": ...}`, or a nested structure. The default collate function knows how to recurse through tuples, lists, dicts and namedtuples, so whatever shape you pick is preserved with an extra leading batch dimension. The base `Dataset` class supplies almost nothing: it defines `__add__`, which returns a `ConcatDataset`, and leaves `__getitem__` for you. There is no hidden machinery — the value is in the convention that `DataLoader`, `Subset`, `random_split` and `ConcatDataset` all rely on. Put the expensive work inside `__getitem__`, not `__init__`. `__init__` should build a lightweight index — a list of file paths, an offsets table, a dataframe of labels — while `__getitem__` does the read, the decode and the transform. That is what gets distributed across worker processes; anything eagerly loaded in `__init__` is loaded once, in the parent, and then duplicated into every worker. ## What DataLoader adds 1. **Order.** `shuffle=True` wraps the dataset in a `RandomSampler`; `shuffle=False` gives a `SequentialSampler`. You can pass your own `sampler` (for example `WeightedRandomSampler` for class balancing), but `sampler` and `shuffle` are mutually exclusive and passing both raises `ValueError`. `batch_sampler` goes one level further and yields whole lists of indices, which is how length-bucketing is implemented. 2. **Batching.** `batch_size` indices are drawn, `__getitem__` is called for each, and the resulting samples are handed to `collate_fn`. 3. **Collation.** The default collate stacks same-shaped tensors along a new dimension 0, converts NumPy arrays and Python numbers to tensors, and recurses through containers. 4. **Parallelism.** `num_workers > 0` moves fetching and collation into separate processes. 5. **Transfer.** `pin_memory=True` copies each batch into page-locked memory so the host-to-device copy can be asynchronous. 6. **Edges.** `drop_last=True` discards a final short batch; `timeout` bounds how long the main process waits on a worker. A `DataLoader` is not an iterator, it is an iterable: each `for batch in loader` creates a fresh iterator, which is why every epoch reshuffles. ## Frequent mistakes **Returning a batch from `__getitem__`.** If your `__getitem__` returns a tensor of shape `(32, 3, 224, 224)` and you also set `batch_size=32`, you get `(32, 32, 3, 224, 224)`. One index, one sample. **Forgetting `__len__`.** Shuffling needs the length; without it `len(dataset)` raises `TypeError` and the default sampler cannot be built. **dtype drift.** The default collate converts a NumPy array with `torch.as_tensor`, preserving dtype. A `np.float64` array becomes a `torch.float64` tensor, and feeding that to a float32 `nn.Linear` raises a scalar-type error. Cast in `__getitem__`. **Non-picklable handles.** An open HDF5 file, an LMDB environment or a database connection created in `__init__` breaks or corrupts under multiprocessing. Open them lazily, on first use inside the worker. **Shared transforms across splits.** `random_split` returns `Subset` views over the *same* parent dataset, so the parent's augmentation pipeline applies to the validation subset too. Build two dataset objects with different transforms instead.

What happens if you pass both shuffle=True and a custom sampler to a DataLoader?
It raises a ValueError — they are mutually exclusive. `shuffle=True` is just a shorthand for installing a `RandomSampler`, so once you supply your own sampler you own the ordering. If you want randomness with a custom sampler, make the sampler itself random (for example `WeightedRandomSampler`, or a `SubsetRandomSampler`) and leave `shuffle` unset.
When would you set drop_last=True?
When a short final batch would hurt. A batch of size 1 makes `nn.BatchNorm` raise in training mode because it cannot compute a variance per channel; a ragged last batch also skews per-batch metric averaging and forces recompilation or reshaping in graph-capture paths like `torch.compile` or CUDA graphs. On validation you usually leave it False so no samples are silently dropped.
How do random_split and Subset relate to a Dataset?
`random_split` returns `Subset` objects, each holding a reference to the same parent dataset plus a list of indices — no data is copied, and indexing a `Subset` just forwards to the parent's `__getitem__`. The consequence people miss is that the parent's transform applies to every split, so a train-time augmentation will also run on your validation subset unless you construct separate dataset objects.

saying these in an interview costs you the question

Says __getitem__ should return a whole batch
Thinks Dataset performs the shuffling
Loads the entire dataset eagerly in __init__
Claims __len__ is optional for a shuffled map-style dataset
Believes DataLoader moves batches to the GPU by itself

In PyTorch, what must a custom map-style Dataset implement, and what does DataLoader add?

§ answer

basics

solid answer

code

§ go deeper

▾level notes4 levels

▾deep dive~3 min

Two dataset styles

The map-style contract

What DataLoader adds

Frequent mistakes

▾common follow-ups3 questions

▾red flags5 flags

§ context

▾related questions5 questions

▾related topics3 topics

Two dataset styles

The map-style contract

What DataLoader adds

Frequent mistakes

answer

go deeper

context