Iterate Through DataLoader

mediumPRODatasetsPyTorch DatasetsPytorch

Problem Statement

Write a function that iterates through a DataLoader and collects batch statistics. This simulates the core loop used during model training.

Background

During training, you iterate through the DataLoader to process data in batches. Each iteration yields one batch, and the last batch may be smaller if the dataset size isn't evenly divisible by batch size.

Your Task

Write a function get_batch_stats(dataloader) that iterates through all batches and returns statistics about them.

Output Format

Return a dictionary with keys: "num_batches" (int), "first_batch_size" (int), "last_batch_size" (int), "total_samples" (int).

Example

Input

7 samples, batch_size=3

Output

{"num_batches": 3, "first_batch_size": 3, "last_batch_size": 1, "total_samples": 7}

Reasoning

Input: a loader over a dataset of $N = 7$ samples (each a feature row plus a label), served in order with batch size $B = 3$ . A batch's size is the number of samples in it.
Batching: samples are taken $B$ at a time until the dataset runs out, so there are $\lceil N/B \rceil = \lceil 7/3 \rceil = 3$ batches. The first two are full ( $3$ and $3$ ); the third gets the leftover $7 - 2 \cdot 3 = 1$ sample, i.e. $7 \bmod 3 = 1$ .
Statistics: num_batches counts the batches ( $3$ ), first_batch_size is the size of batch 1 ( $3$ ), last_batch_size is the size of the final batch ( $1$ ), and total_samples adds all sizes: $3 + 3 + 1 = 7$ .
Notice: total_samples always recovers $N$ , while last_batch_size equals $B$ only when $B$ divides $N$ (6 samples with $B = 2$ gives $2, 2, 2$ ). If $B > N$ , one batch holds everything, so first and last are the same batch (5 samples with $B = 10$ gives $5$ and $5$ ).

Constraints

DataLoader will have at least one batch
Features and labels are the first two elements returned
Last batch may be smaller than batch_size

Learn this in

PyTorch FundamentalsFoundations

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy

Concept primerPREVIEW

1. Background Knowledge

PyTorch DataLoader is a core utility for efficient data batching during ML training. It wraps a Dataset and provides iterable batches via drop_last=False (default), where the final batch may be smaller if the dataset size $N$ is not perfectly divisible by batch_size $B$ : $\text{num_batches} = \lceil N/B \rceil$ .

Key concepts:

Batching: Divides dataset into subsets of size $B$ (except possibly last batch).
Iteration: for batch_features, batch_labels in dataloader: yields tuples where len(batch_features) == batch_size (or smaller for last).
Total samples: $\sum$ \text{batch sizes across all batches} $= N$ .
Prerequisites: Familiarity with PyTorch Dataset, DataLoader, and Python iteration/ unpacking.

2. Algorithm Approach

Use a single-pass linear scan over the DataLoader iterator:

Initialize counters: num_batches = 0, total_samples = 0.
Track first/last batch sizes via flags or list.
For each batch: increment counters, get batch_size = len(batch_features).
No data storage needed—compute statistics on-the-fly.

This mirrors training loops but collects metadata instead of training.

3. Step-by-Step Strategy

def get_batch_stats(dataloader):
 stats = {
 "num_batches": 0,
 "first_batch_size": 0,
 "last_batch_size": 0,
 "total_samples": 0
 }

 first_batch = True
 for batch_features, batch_labels in dataloader:
 batch_size = len(batch_features) # Assumes features first, valid per constraints

 stats["num_batches"] += 1
 stats["total_samples"] += batch_size

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.