Pytorch
Concatenate Datasets
Dataloader
Machine Learning
Data Preprocessing

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.

Practice ML system design

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.

python
1from torch.utils.data import Dataset, ConcatDataset, DataLoader
2
3class Numbers(Dataset):
4    def __init__(self, values):
5        self.values = values
6
7    def __len__(self):
8        return len(self.values)
9
10    def __getitem__(self, idx):
11        x = self.values[idx]
12        return x, x * 10
13
14
15d1 = Numbers([1, 2, 3])
16d2 = Numbers([4, 5])
17combined = ConcatDataset([d1, d2])
18loader = DataLoader(combined, batch_size=2, shuffle=True)
19
20for batch in loader:
21    print(batch)

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.

python
print(d1[0])
print(d2[0])

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.

python
1for epoch in range(2):
2    for features, targets in loader:
3        # forward pass, loss, backward pass, optimizer step
4        pass

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 ConcatDataset balances 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 ConcatDataset would have been simpler and more memory-efficient.

Summary

  • In PyTorch, use ConcatDataset to combine datasets before creating a DataLoader.
  • 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
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track 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.

Practice ML system design

All Rights Reserved.