Datasets and DataLoaders: Worked examples — AI-ML deep dive

This is the drill-yard for Datasets and DataLoaders. The parent note built the machinery: __len__, __getitem__, batching, collation, drop_last. Here we push that machinery through every corner case — the ugly divisions, the empty inputs, the ragged batches — so that when your own DataLoader misbehaves at 2 a.m., you have already seen the failure.

Before we count a single batch, let us re-earn the one symbol everything depends on.

We need one tool to answer that: integer division with a remainder. Not ordinary division — because you cannot train on "half a handful."

The drop_last idea shows up the moment your batches feed mini-batch gradient descent and again in the training loop, so it is worth beating flat.

The scenario matrix

Every question this topic can throw at you is one of these cells. The worked examples below are tagged with the cell(s) they cover.

Cell	Situation	What's tricky	Example
A	$N$ divides evenly by $B$ ( $N mod B = 0$ )	floor = ceiling, no leftover	Ex 1
B	$N$ does not divide evenly, `drop_last=False`	ragged last batch survives	Ex 2
C	Same, `drop_last=True`	you lose samples every epoch	Ex 3
D	Degenerate: $B \geq N$ (batch bigger than dataset)	exactly one batch, maybe short	Ex 4
E	Degenerate: $B = 1$	$N$ batches, pure stochastic; and $N = 0$ empty jar	Ex 5
F	Collation of ragged samples (variable length)	must pad to a rectangle	Ex 6
G	Real-world word problem: epochs, steps, wall-clock	chain the formulas together	Ex 7
H	Exam twist: `shuffle` + multi-epoch sample counting	shuffling changes order not count	Ex 8

The worked examples

Worked example Example 1 — Cell A: clean division

A Dataset has $N = 128$ images. You set batch_size=32, drop_last=False. How many batches per epoch, and how big is the last one?

Forecast: Guess before reading. Does $32$ divide $128$ ? If yes, no leftover — write down your number now.

Steps

Compute the remainder: $128 mod 32 = 0$ . Why this step? The remainder is the single fact that decides whether there is a ragged batch. Zero remainder ⇒ every batch is full.
Because remainder is $0$ , floor and ceiling agree: $⌊ 128/32 ⌋ = ⌈ 128/32 ⌉ = 4$ . Why this step? When nothing is left over, "complete handfuls" and "handfuls including scrap" are the same count — drop_last makes no difference here.
Last batch size $= B = 32$ (the $N mod B = 0$ branch of the formula). Why this step? No scrap means the final scoop is also a full $32$ .

Answer: $4$ batches, each of size $32$ .

Verify: $4 \times 32 = 128 = N$ . Every sample used exactly once, none dropped. ✓

Worked example Example 3 — Cell C: same data,

drop_last=True Same $N = 100$ , batch_size=32, but now drop_last=True. Batches per epoch? How many samples are never seen this epoch?

Forecast: You already know 4 samples are the tail. What happens to them now?

Steps

Batches $= ⌊ 100/32 ⌋ = 3$ . Why this step? drop_last=True throws away any incomplete final handful, so we use floor, not ceiling.
Samples used $= 3 \times 32 = 96$ . Why this step? Only full batches survive.
Samples dropped this epoch $= 100 - 96 = 4$ (which is exactly $N mod B$ ). Why this step? The dropped count is always the remainder — it is the tail we refused to keep.

Answer: $3$ batches; $4$ samples skipped this epoch.

Verify: Dropped $= N mod B = 100 mod 32 = 4$ . ✓

[!mistake] drop_last=True is not "lose data forever" Because shuffle=True reshuffles every epoch, a different 4 samples get dropped next epoch. Over many epochs the loss is statistically tiny — but for a tiny validation set it can bite. Never use drop_last=True on evaluation loaders.

Worked example Example 4 — Cell D: batch bigger than the dataset

A quick smoke-test Dataset has $N = 5$ samples. You forgot to shrink the batch size: batch_size=64, drop_last=False. What comes out?

Forecast: Can you scoop 64 marbles from a jar of 5?

Steps

Full handfuls: $⌊ 5/64 ⌋ = 0$ . You cannot fill even one complete batch. Why this step? A scoop only counts as "full" if it reaches $B$ ; here it never does.
Ceiling: $⌈ 5/64 ⌉ = 1$ . drop_last=False still yields the one short batch. Why this step? All 5 samples exist and are handed over together — one ragged batch of size 5.
If instead drop_last=True: $⌊ 5/64 ⌋ = 0$ ⇒ zero batches, the loop body never runs. Why this step? With no complete batch and dropping enabled, the epoch is silently empty — a classic "why is my model not training?" bug.

Answer: drop_last=False → 1 batch of size 5. drop_last=True → 0 batches.

Verify: $0 \times 64 = 0$ used when dropping; $1$ batch $\times 5 = 5 = N$ when keeping. ✓

Worked example Example 5 — Cell E: the extremes

$B = 1$ and $N = 0$ (a) $N = 7$ , batch_size=1. Batches? (b) An empty Dataset $N = 0$ , any $B$ . Batches?

Forecast: Batch size 1 means every scoop is a single marble — how many scoops? And an empty jar?

Steps

(a) $⌈ 7/1 ⌉ = ⌊ 7/1 ⌋ = 7$ . drop_last is irrelevant since $7 mod 1 = 0$ . Why this step? $B = 1$ is pure online / stochastic gradient descent — one sample per step, $N$ steps.
(a) Each "batch" tensor still has a leading batch dimension of size 1, e.g. shape (1, 3, 224, 224). Why this step? The DataLoader always adds a batch axis, even for one sample — models expect it.
(b) $⌊ 0/ B ⌋ = ⌈ 0/ B ⌉ = 0$ for any $B \geq 1$ . Why this step? No samples ⇒ no batches ⇒ the for batch in loader loop body runs zero times. No crash, just silence.

Answer: (a) $7$ batches of size 1. (b) $0$ batches.

Verify: (a) $7 \times 1 = 7 = N$ ✓. (b) $0 = N$ ✓.

Worked example Example 6 — Cell F: collating ragged sequences (geometric)

A batch of $3$ token sequences of lengths $4$ , $2$ , $5$ , each token a vector of dimension $d = 3$ . Default torch.stack fails — why, and what does the padded batch tensor's shape become?

Forecast: What single number must all three sequences agree on before they can be stacked into one rectangular tensor?

Steps

Default collation calls torch.stack, which needs every sample to have identical shape. Here shapes are $(4, 3), (2, 3), (5, 3)$ — the length axis disagrees. Why this step? A tensor is a rectangle (a grid); you cannot stack rows of different width into one grid. Look at the figure: the three chalk rows have different lengths.
Find $max_len = max (4, 2, 5) = 5$ . Why this step? We must grow every sequence to the longest so they line up — that is the target width of the rectangle.
Pad each shorter sequence with zeros on the length axis: add $1$ row to the length-4, add $3$ rows to the length-2, add $0$ to the length-5 (the pink zero-cells in the figure). Why this step? Padding on the length axis (dim 0), never the feature axis (dim 1) — the feature dimension $d = 3$ must stay meaningful.
Stack the three now-identical $(5, 3)$ tensors ⇒ batch shape $(3, 5, 3)$ , i.e. $(batch, max_len, d)$ . Why this step? Now all shapes match, so torch.stack builds the batch cube.

Answer: batch tensor shape $(3, 5, 3)$ ; total padded cells added $= (5 - 4) + (5 - 2) + (5 - 5) = 1 + 3 + 0 = 4$ .

Verify: Real cells $= 4 + 2 + 5 = 11$ tokens; padded total $= 3 \times 5 = 15$ ; padding added $= 15 - 11 = 4$ . ✓ (Matches step 3.) You'd then hand the model a mask so it ignores those 4 pad cells.

Worked example Example 7 — Cell G: real-world word problem (epochs → wall-clock)

You have $N = 50, 000$ images, batch_size=64, drop_last=True, training for $8$ epochs. Each optimizer step (one batch) takes $0.05$ seconds. (a) Steps per epoch? (b) Total steps? (c) Total training time?

Forecast: "Steps" and "batches" are the same thing here — estimate the total before computing.

Steps

Steps per epoch $= ⌊ 50000/64 ⌋ = 781$ (since drop_last=True ⇒ floor). Check: $781 \times 64 = 49984$ , remainder $16$ dropped. Why this step? One optimizer update per full batch; the ragged 16 are dropped, so use floor.
Total steps $= 8 \times 781 = 6248$ . Why this step? Each epoch replays the (reshuffled) data, giving the same step count each time.
Total time $= 6248 \times 0.05 = 312.4$ seconds $\approx 5.2$ minutes. Why this step? Time = steps × seconds-per-step. Units: $[steps] \times [s/step] = [s]$ . ✓

Answer: (a) $781$ steps/epoch, (b) $6248$ total steps, (c) $312.4$ s.

Verify: $8 \times 781 = 6248$ ; $6248 \times 0.05 = 312.4$ ; dropped per epoch $= 50000 mod 64 = 16$ . ✓ (This is exactly the throughput reasoning that matters in distributed training, where you split these steps across GPUs.)

Worked example Example 8 — Cell H: exam twist, shuffle across epochs

$N = 10$ , batch_size=4, shuffle=True, drop_last=False, run for $3$ epochs. (a) Over the whole run, how many times is each sample fed to the model? (b) How many total batches are yielded across all $3$ epochs?

Forecast: Does shuffle=True change how many times a sample appears, or only the order?

Steps

Per epoch, shuffle permutes indices but every index still appears exactly once — shuffling reorders, it does not duplicate or drop. Why this step? This is the trap. Randomizing order ≠ random sampling with replacement; each epoch is a full pass.
So each sample is seen once per epoch ⇒ over $3$ epochs, each sample is fed $3$ times. Why this step? "Epochs" literally means "number of full passes over the data."
Batches per epoch: $⌈ 10/4 ⌉ = 3$ (two batches of 4, one ragged batch of $10 mod 4 = 2$ ). Total $= 3 \times 3 = 9$ . Why this step? drop_last=False keeps the size-2 tail; multiply by epochs.

Answer: (a) each sample fed $3$ times; (b) $9$ batches total.

Verify: Per epoch $3 \times$ : sizes $4 + 4 + 2 = 10 = N$ ✓. Across 3 epochs, $9$ batches; total samples processed $= 9 batches \to 3 \times 10 = 30 = 3 N$ sample-visits ✓.

Recall Quick self-test

Batches when $N = 100, B = 32, drop_last=False$ ? ::: $⌈ 100/32 ⌉ = 4$ (last one size 4). Batches when $N = 100, B = 32, drop_last=True$ ? ::: $⌊ 100/32 ⌋ = 3$ ; 4 samples dropped. $B = 64, N = 5, drop_last=True$ gives how many batches? ::: $0$ — the loop never runs. Padded shape of a batch of 3 sequences (lengths 4,2,5, feature dim 3)? ::: $(3, 5, 3)$ . Does shuffle=True change how many times a sample is seen per epoch? ::: No — exactly once; only the order changes. Steps/epoch for $N = 50000, B = 64, drop_last=True$ ? ::: $⌊ 50000/64 ⌋ = 781$ .

Next: plug these batches into the training loop, recall why batching stabilizes mini-batch gradient descent, and see how the same counting logic scales in distributed training and reuse pretrained pipelines in transfer learning.