How do you alter the size of a Pytorch Dataset?

PyTorch

Dataset

Machine Learning

Data Manipulation

Python

How do you alter the size of a Pytorch Dataset?

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

PyTorch Machine Learning Python

Introduction

PyTorch datasets are sized by the __len__ method and indexed by __getitem__. To alter a dataset's size, you can use torch.utils.data.Subset to select specific indices, torch.utils.data.random_split to split into smaller datasets, or override __len__ in a custom wrapper. These approaches avoid copying data — they create views over the original dataset.

Using Subset (Select Specific Indices)

Subset takes a dataset and a list of indices, creating a smaller dataset without copying data:

python

1import torch
2from torch.utils.data import Dataset, Subset, DataLoader
3
4# Original dataset with 10,000 samples
5dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True)
6print(len(dataset))  # 10000 (using subset of MNIST for example)
7
8# Take only the first 1,000 samples
9small_dataset = Subset(dataset, indices=range(1000))
10print(len(small_dataset))  # 1000
11
12# Take random 500 samples
13indices = torch.randperm(len(dataset))[:500].tolist()
14random_subset = Subset(dataset, indices=indices)
15print(len(random_subset))  # 500
16
17# Use with DataLoader as normal
18loader = DataLoader(random_subset, batch_size=32, shuffle=True)

Using random_split

random_split divides a dataset into non-overlapping subsets:

python

1from torch.utils.data import random_split
2
3dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)
4print(len(dataset))  # 50000
5
6# Split into train (40000), validation (5000), unused (5000)
7train_set, val_set, _ = random_split(dataset, [40000, 5000, 5000])
8print(len(train_set))  # 40000
9print(len(val_set))    # 5000
10
11# With a generator for reproducibility
12generator = torch.Generator().manual_seed(42)
13train_set, val_set = random_split(dataset, [45000, 5000], generator=generator)

Custom Dataset Wrapper (Override len)

For more control, wrap the dataset and override __len__ and __getitem__:

python

1from torch.utils.data import Dataset
2
3class ResizedDataset(Dataset):
4    def __init__(self, dataset, size):
5        self.dataset = dataset
6        self.size = min(size, len(dataset))
7
8    def __len__(self):
9        return self.size
10
11    def __getitem__(self, idx):
12        # Map index to the original dataset
13        return self.dataset[idx]
14
15# Limit dataset to 2000 samples
16original = torchvision.datasets.MNIST(root='./data', train=True)
17smaller = ResizedDataset(original, size=2000)
18print(len(smaller))  # 2000

Repeating / Expanding a Dataset

To make a dataset appear larger (useful for balancing classes or matching epoch sizes):

python

1class RepeatedDataset(Dataset):
2    def __init__(self, dataset, target_size):
3        self.dataset = dataset
4        self.target_size = target_size
5
6    def __len__(self):
7        return self.target_size
8
9    def __getitem__(self, idx):
10        # Wrap around using modulo
11        return self.dataset[idx % len(self.dataset)]
12
13# Small dataset repeated to appear larger
14small_data = [1, 2, 3, 4, 5]  # 5 items
15expanded = RepeatedDataset(small_data, target_size=20)
16print(len(expanded))  # 20
17print(expanded[7])    # 3 (index 7 % 5 = 2, small_data[2] = 3)

ConcatDataset (Combine Multiple Datasets)

python

1from torch.utils.data import ConcatDataset
2
3dataset_a = torchvision.datasets.MNIST(root='./data', train=True)
4dataset_b = torchvision.datasets.MNIST(root='./data', train=False)
5
6combined = ConcatDataset([dataset_a, dataset_b])
7print(len(combined))  # 60000 + 10000 = 70000

Filtering by Condition

Select samples that match a condition (e.g., specific classes):

python

1class FilteredDataset(Dataset):
2    def __init__(self, dataset, condition):
3        self.dataset = dataset
4        self.indices = [i for i in range(len(dataset)) if condition(dataset[i])]
5
6    def __len__(self):
7        return len(self.indices)
8
9    def __getitem__(self, idx):
10        return self.dataset[self.indices[idx]]
11
12# Keep only digits 0-4 from MNIST
13mnist = torchvision.datasets.MNIST(root='./data', train=True, download=True)
14small_digits = FilteredDataset(mnist, condition=lambda x: x[1] < 5)
15print(len(small_digits))  # ~30000 (roughly half)

Or use Subset with filtered indices:

python

indices = [i for i, (_, label) in enumerate(mnist) if label < 5]
small_digits = Subset(mnist, indices)

Using Samplers for Effective Size Control

Instead of modifying the dataset, use a sampler with the DataLoader:

python

1from torch.utils.data import DataLoader, RandomSampler
2
3dataset = torchvision.datasets.MNIST(root='./data', train=True)
4
5# Only sample 1000 items per epoch (without altering the dataset)
6sampler = RandomSampler(dataset, num_samples=1000, replacement=False)
7loader = DataLoader(dataset, batch_size=32, sampler=sampler)
8
9# Each epoch iterates over 1000 samples
10for batch in loader:
11    pass  # 1000 / 32 ≈ 32 batches per epoch

Common Pitfalls

Modifying __len__ without adjusting __getitem__: If __len__ returns 100 but __getitem__ can access indices beyond 100, the DataLoader works fine — but manual iteration with for i in range(len(dataset)) will miss samples. Keep them consistent.
Using random_split without a seed: Without a Generator seed, random_split produces different splits each run. For reproducible experiments, always pass generator=torch.Generator().manual_seed(42).
Large index lists consuming memory: Storing 10 million indices in a Python list uses significant memory. For very large datasets, use a range-based approach or a sampler instead of materializing all indices.
Forgetting that Subset does not copy data: Subset holds a reference to the original dataset. If the original is modified or deleted, the subset breaks. Keep the original dataset alive for the subset's lifetime.
Shuffle conflicts with custom samplers: DataLoader(shuffle=True) and sampler=... are mutually exclusive. When using a RandomSampler or other custom sampler, do not set shuffle=True.

Summary

Use Subset(dataset, indices) to select specific samples by index
Use random_split(dataset, [sizes]) to split into non-overlapping subsets
Override __len__ and __getitem__ in a wrapper class for custom sizing logic
Use ConcatDataset to combine multiple datasets into one
Use RandomSampler(num_samples=N) to limit samples per epoch without altering the dataset
Always set a random seed with torch.Generator().manual_seed() for reproducible splits

Tackling System Design Interview Problems

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

Start the free course

Basis Lab

No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my 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

How do you alter the size of a Pytorch Dataset?

ML System Design practice on Codemia

Introduction

Using Subset (Select Specific Indices)

Using random_split

Custom Dataset Wrapper (Override len)

Repeating / Expanding a Dataset

ConcatDataset (Combine Multiple Datasets)

Filtering by Condition

Using Samplers for Effective Size Control

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

No course covers this one yet

Track what you have practised

ML System Design practice on Codemia