Pytorch - Concatenating Datasets before using Dataloader
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
In PyTorch, the normal way to combine multiple datasets before passing them to a DataLoader is torch.utils.data.ConcatDataset. It gives you a single dataset interface over several underlying datasets without copying all samples into one large in-memory structure.
That is usually the right answer when the datasets have the same sample format and should be treated as one training source.
Use ConcatDataset for Sequential Composition
ConcatDataset takes a list of datasets and exposes them as one dataset whose length is the sum of the parts.
The loader treats combined like any other dataset.
What ConcatDataset Actually Does
ConcatDataset does not merge samples physically. It stores references to the original datasets and maps a global index to the correct child dataset and local index.
That means:
- memory use stays reasonable
- updates to the underlying datasets are reflected if those datasets are dynamic
- all child datasets must return compatible sample structures for batching
The last point is the one that breaks most often.
Make Sure Sample Shapes Match
If one dataset returns (image, label) and another returns only image, the DataLoader collate step will fail or behave inconsistently.
A good sanity check is to inspect one sample from each dataset before concatenating.
If the data types, tensor shapes, or target formats differ, standardize them first.
When You Need Different Sampling Behavior
Concatenation is not always the full solution. Sometimes you also need to control class balance or dataset mixing frequency.
For example, if dataset A has 100000 samples and dataset B has 500 samples, plain concatenation means B will appear rarely. In that case, consider:
- weighted sampling
- oversampling the smaller dataset
- custom samplers
- separate loaders with explicit training logic
ConcatDataset solves composition, not balancing.
Concatenation Versus Zipping
Do not confuse concatenating with pairing. ConcatDataset appends datasets one after another. If you need sample i from dataset A aligned with sample i from dataset B, that is a different problem and usually needs a custom Dataset implementation.
A Practical Training Example
Once concatenated, training code stays unchanged.
That is why ConcatDataset is so useful: it keeps the combination step separate from the rest of the training loop.
Common Pitfalls
- Concatenating datasets that return incompatible sample shapes or target formats.
- Assuming
ConcatDatasetbalances datasets automatically when one is much larger than another. - Using concatenation when the real need is aligned pairing, not sequential appending.
- Forgetting that transforms should be compatible across all child datasets.
- Rebuilding a huge merged dataset manually when
ConcatDatasetwould have been simpler and more memory-efficient.
Summary
- In PyTorch, use
ConcatDatasetto combine datasets before creating aDataLoader. - It provides one logical dataset without copying everything into memory.
- All child datasets must return compatible sample structures.
- Concatenation changes composition, not class balance or sampling policy.
- If you need aligned pairs instead of appended samples, build a custom dataset instead.
Related reading
- Pytorch - Concatenating Datasets before using Dataloader
- pytorch - connection between loss.backward and optimizer.step
- PyTorch - How to get learning rate during training?
- Pytorch - Using more GPUs and increasing batch size makes training slower in DistributedDataParallel
- Pytorch Autograd what does runtime error grad can be implicitly created only for scalar outputs mean
- PyTorch Binary Classification - same network structure, ''simpler'' data, but worse performance?
- quadratic featurizer preprocessing with fit_transform
- Query that will find users who post THE SAME SET of marks as user2
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.