Iterate Through DataLoader
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:
7 samples, batch_size=3
{"num_batches": 3, "first_batch_size": 3, "last_batch_size": 1, "total_samples": 7}7 samples with batch_size=3 gives batches of sizes [3, 3, 1]
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
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: ∑\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
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.