Skip to main content

Batching from scratch, then DataLoader

Shuffle the indices with a seeded generator, cut them into batches, keep or drop the short last one, and stack the items: a loader in fifteen lines. Then DataLoader does the same job, batch for batch, and its default collate function is opened up.
DL01.9.2L2 Intermediate25 min
In review

The maths and code on this page are checked by automated tests on every change, but an independent reader has not signed it off yet. Spot a mistake or a confusing step? Tell us.

Before this lesson: DL01.9.1 A dataset is something you can index

What you will be able to do
  • Write a batch generator that shuffles with a seeded generator, cuts batches of a given size and stacks the items.
  • Count the steps of an epoch, ceil(n/|B|) or floor(n/|B|) with drop_last, and the size of the last batch, for the 1,077 training images.
  • Reproduce the batches of the scratch generator with DataLoader given the same permutation, and make a shuffled DataLoader repeatable with a seeded generator.
  • Explain what the default collate function does to a list of (input, label) pairs.

The question

The 1,077 training images of this course go to the network in batches of 64. How many steps make one epoch, what happens to the 53 images that do not fill a last batch, and how does a loader decide which images travel together?

From the last lesson

1. Which two methods make an object a map-style dataset?

Show answer

__len__, returning the number of items, and __getitem__(i), returning item \(i\).

2. What shape and dtype does the input of one digits item have?

Show answer

Shape (64,), the 64 pixels divided by 16, and dtype float32.

Intuition

A loader never looks inside the data. It works on the indices \(0, 1, \dots, {n - 1}\) and does three things with them in every epoch: it shuffles them, cuts the shuffled list into runs of \(\abs{B}\), and for each run fetches the items from the dataset and stacks them into one tensor.

Four rows built from ten numbered tiles. Top row, the dataset: tiles 0 to 9 in order. An arrow labelled randperm(10, generator=g) leads to row 1, shuffle: 4, 1, 7, 5, 3, 9, 0, 8, 6, 2. Row 2, cut with batch size 4: the same order split into three groups, 4 1 7 5 tinted blue, 3 9 0 8 tinted sky blue, and 6 2 with dashed edges marked short, labelled batch 1, batch 2 and batch 3. Row 3, stack: under each group a box with its tensor shapes, X (4, 64) and y (4,) for the first two, X (2, 64) and y (2,) for the third. A note says drop_last=True would throw batch 3 away: 2 steps instead of 3.
DL01.9.2-F1

One epoch of a loader on ten examples with \(\abs{B} = 4\): shuffle the indices, cut them into batches, stack each batch. The last batch is short; drop_last decides whether it is used.

An epoch is one pass over the data, so each index appears in exactly one batch. A new epoch shuffles again, with the generator moved on, so the batches are different every time.

Build it up

Step 1. Shuffle. torch.randperm(n, generator=g) returns the numbers \(0\) to \({n - 1}\) in a random order. The generator g is a random-number source with its own state: seeded with torch.Generator().manual_seed(0), it gives the same sequence of permutations in every run, and each call moves its state on, so the next epoch gets a new order. With \(n = 10\) the first call gives \(4, 1, 7, 5, 3, 9, 0, 8, 6, 2\).

Step 2. Cut. Take the shuffled list in runs of \(\abs{B}\): positions 0 to 3, then 4 to 7, and so on. With \(\abs{B} = 4\) and \(n = 10\) that is two full batches and a short one of two.

Step 3. Keep or drop the short batch. With drop_last=False, the default, the short batch is used as a smaller step. With drop_last=True it is skipped for this epoch.

Step 4. Stack. For each batch, fetch the items with ds[i] and join them. torch.stack puts \(\abs{B}\) tensors of shape (64,) on top of one another as the rows of one tensor of shape (|B|, 64); the labels, scalar tensors, become a vector of shape (|B|,). In DataLoader this job belongs to the collate function.

Left, a list of four items numbered 0 to 3, each a blue box x of shape (64,) beside a small black box y. An arrow labelled default_collate leads right to a list of two tensors: a tall blue block of four rows labelled X: (4, 64) float32, and a column of four small black cells labelled y: (4,). A note says a new first axis, the batch axis, is added; the tuple becomes a list.
DL01.9.2-F3

The default collate function. A list of \(\abs{B}\) pairs becomes a list of two tensors, each with a new first axis of length \(\abs{B}\).

Step 5. Repeat. An outer loop over epochs calls the batch generator again; DL01.8 wrote that loop, and this lesson only changes where its batches come from.

The formal statement

Divide the number of examples \(n\) by the batch size with remainder:

\[ n = q\,\abs{B} + r, \qquad 0 \le r < \abs{B}. \]

Here \(q = \lfloor n / \abs{B} \rfloor\) is the number of full batches and \(r\) the number left over.

  • If \(r = 0\), every batch is full, and an epoch has \(q\) steps whether or not the last is dropped.
  • If \(r > 0\), drop_last=False gives \(q\) full batches and one more of size \(r\), so \({q + 1}\) steps; drop_last=True gives \(q\) steps and leaves \(r\) examples out of this epoch.

Both cases fit two formulas. Without dropping, an epoch has

\[ K = \left\lceil \frac{n}{\abs{B}} \right\rceil \]

steps, because \(\lceil n / \abs{B} \rceil = q\) when \(r = 0\) and \({q + 1}\) when \(r > 0\). With drop_last=True it has \(\lfloor n / \abs{B} \rfloor = q\) steps.

Which examples are dropped. With shuffling, the \(r\) examples left over are those in the last \(r\) places of a random permutation. Every place is equally likely for every example, so each example is left out of a given epoch with probability \(r / n\), and a different set is left out each epoch.

Three rows labelled epoch 1, epoch 2 and epoch 3, each showing the ten indices 0 to 9 in the order a DataLoader with batch_size 4, shuffle True and drop_last True gives them: 3 7 5 2 | 0 8 1 6 | 9 4, then 1 3 2 7 | 5 9 0 6 | 4 8, then 6 8 5 7 | 2 4 1 0 | 3 9. The first two groups in each row are tinted blue and sky blue; the last two tiles have grey dashed edges and the label dropped. Indices 9 and 4 are dropped in epoch 1, 4 and 8 in epoch 2, 3 and 9 in epoch 3.
DL01.9.2-F4

A fresh shuffle every epoch. With drop_last=True two of the ten examples sit out each epoch, but not the same two, so over several epochs every example is used.

What DataLoader builds. DataLoader(ds, batch_size, shuffle, drop_last, generator) puts the steps above together from three parts:

Part Its job
sampler the order of the indices
batch sampler cutting them into lists of \(\abs{B}\)
collate function stacking the items

shuffle=False gives a sequential sampler, \(0, 1, \dots, {n - 1}\). shuffle=True gives a RandomSampler, which draws a new permutation each epoch from generator. You may instead pass any sampler, for example a list of indices, but then shuffle must not be set. The default collate function adds a first axis to tensors, turns NumPy arrays and Python numbers into tensors, and keeps the structure of the items, except that a tuple comes back as a list.

Worked example

The 1,077 training images. DL01.9.4 fixes the split; the training part has \(n = 1077\).

Three long bars, each standing for the 1,077 training images cut into batches drawn as alternating blue and sky-blue segments, with the short last batch yellow and dashed. Top, batch size 32: 1077 = 33 times 32 plus 21; steps ceil 34, with drop_last 33; last batch 21. Middle, batch size 64: 1077 = 16 times 64 plus 53; steps ceil 17, with drop_last 16; last batch 53. Bottom, batch size 512: 1077 = 2 times 512 plus 53; steps ceil 3, with drop_last 2; last batch 53.
DL01.9.2-F2

The same 1,077 images cut three ways. The yellow piece is the short last batch.

Batches of 64. \({16 \cdot 64} = 1024\) and \({1077 - 1024} = 53\), so \(q = 16\) and \(r = 53\). An epoch has \(\lceil 1077/64 \rceil = 17\) steps, the last on 53 images, or 16 with drop_last=True. Then 53 images sit out each epoch; any one image does so with probability \({53 / 1077} \approx 0.049\).

Batches of 32. \({33 \cdot 32} = 1056\), leaving 21: 34 steps, or 33.

Batches of 512. \({2 \cdot 512} = 1024\), leaving 53 again: 3 steps, or 2. Dropping the last batch now throws away about one twentieth of the data in every epoch for a saving of one step in three.

Ten examples by hand. With \(n = 10\), \(\abs{B} = 4\) and the permutation of Step 1, \(4, 1, 7, 5, 3, 9, 0, 8, 6, 2\), the batches are \(\{4, 1, 7, 5\}\), \(\{3, 9, 0, 8\}\) and \(\{6, 2\}\): \(\lceil 10/4 \rceil = 3\) steps, or 2 with drop_last=True, which leaves out examples 6 and 2.

Code

The data: the training images of the course’s split (made here with the two calls DL01.9.4 explains), as a TensorDataset.

import math
import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset, default_collate
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
digits = load_digits()
X, y = digits.images.reshape(-1, 64) / 16, digits.target
X_rest, X_test, y_rest, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
X_train, X_val, y_train, y_val = train_test_split(X_rest, y_rest, test_size=0.25, stratify=y_rest,
                                                  random_state=0)
train_set = TensorDataset(torch.tensor(X_train, dtype=torch.float32), torch.tensor(y_train))
print(len(train_set))
1077

The loader from scratch. A Python generator function: each call is one epoch.

def batches(ds, batch_size, shuffle=False, drop_last=False, generator=None, order=None):
    """Yield the batches of one epoch, each a tuple of stacked tensors, such as (X, y)."""
    n = len(ds)
    if order is None:
        order = torch.randperm(n, generator=generator).tolist() if shuffle else list(range(n))
    for start in range(0, n, batch_size):
        idx = order[start:start + batch_size]
        if drop_last and len(idx) < batch_size:
            break
        items = [ds[i] for i in idx]
        yield tuple(torch.stack(field) for field in zip(*items))   # one stacked tensor per field

for size in (1, 32, 64, 512):
    keep = sum(1 for _ in batches(train_set, size))
    drop = sum(1 for _ in batches(train_set, size, drop_last=True))
    print(f"|B| = {size:3}: {keep:4} steps, {drop:4} with drop_last, last batch {len(train_set) - (keep - 1) * size}")
|B| =   1: 1077 steps, 1077 with drop_last, last batch 1
|B| =  32:   34 steps,   33 with drop_last, last batch 21
|B| =  64:   17 steps,   16 with drop_last, last batch 53
|B| = 512:    3 steps,    2 with drop_last, last batch 53

The permutation of the worked example, and one shuffled epoch of the ten-example toy.

perm10 = torch.randperm(10, generator=torch.Generator().manual_seed(0)).tolist()
toy = TensorDataset(torch.arange(10))
print(perm10, [b[0].tolist() for b in batches(toy, 4, order=perm10)])
[4, 1, 7, 5, 3, 9, 0, 8, 6, 2] [[4, 1, 7, 5], [3, 9, 0, 8], [6, 2]]

Now DataLoader. Given the same permutation as its sampler (a list of indices is a valid sampler), it makes exactly the batches of the scratch version.

order = torch.randperm(len(train_set), generator=torch.Generator().manual_seed(0)).tolist()
mine = list(batches(train_set, 64, order=order))
theirs = list(DataLoader(train_set, batch_size=64, sampler=order))
equal = len(mine) == len(theirs) and all(torch.equal(a[0], b[0]) and torch.equal(a[1], b[1])
                                         for a, b in zip(mine, theirs))
print(len(theirs), "batches, all equal:", equal, "  dtypes:", theirs[0][0].dtype, theirs[0][1].dtype)
17 batches, all equal: True   dtypes: torch.float32 torch.int64

With shuffle=True and a seeded generator, a loader is repeatable: two loaders built the same way give the same batches, epoch after epoch, and each epoch differs from the last.

def epochs(seed, count=3, n=10):
    loader = DataLoader(TensorDataset(torch.arange(n)), batch_size=4, shuffle=True, drop_last=True,
                        generator=torch.Generator().manual_seed(seed))
    return [[b[0].tolist() for b in loader] for _ in range(count)]

first, again = epochs(0), epochs(0)
print(first)
print("repeatable:", first == again, "  dropped each epoch:", [sorted(set(range(10)) - set(sum(e, []))) for e in first])
[[[3, 7, 5, 2], [0, 8, 1, 6]], [[1, 3, 2, 7], [5, 9, 0, 6]], [[6, 8, 5, 7], [2, 4, 1, 0]]]
repeatable: True   dropped each epoch: [[4, 9], [4, 8], [3, 9]]

Finally, the collate function on its own: a list of pairs in, a list of two stacked tensors out.

pairs = [train_set[i] for i in range(4)]
X4, y4 = default_collate(pairs)
print(type(default_collate(pairs)).__name__, X4.shape, y4.shape, torch.equal(X4, torch.stack([p[0] for p in pairs])))
list torch.Size([4, 64]) torch.Size([4]) True
Go deeper: why a shuffled DataLoader and a bare randperm with the same seed disagree

In PyTorch 2.15, when a loader starts an epoch, its iterator first draws one number from generator, a base seed for worker processes, and only then does the RandomSampler call torch.randperm(n, generator=generator). After the permutation is used up, the sampler makes one more, empty-sliced call to randperm. So with manual_seed(0) the loader’s first epoch is the second thing the generator produces, not the first. That is why the code above compares the two versions on the same permutation. The behaviour is an internal detail that may change between versions; what is promised is that the same seed gives the same batches.

g = torch.Generator().manual_seed(0)
torch.empty((), dtype=torch.int64).random_(generator=g)        # the base seed
manual = torch.randperm(10, generator=g).tolist()
loader = DataLoader(TensorDataset(torch.arange(10)), batch_size=4, shuffle=True,
                    generator=torch.Generator().manual_seed(0))
first_epoch = sum((b[0].tolist() for b in loader), [])
print(manual, first_epoch)
[3, 7, 5, 2, 0, 8, 1, 6, 9, 4] [3, 7, 5, 2, 0, 8, 1, 6, 9, 4]

Explore

The epoch shuffler in DL01.9.3 runs this lesson’s loader on 24 tiles. Set the batch size to 5, step through an epoch with shuffle on and drop_last on, and predict before each epoch how many tiles will sit out.

Pitfalls

Two grids with one row per batch, 17 batches of 64, and one column per digit class 0 to 9; a cell is shaded darker the more images of that class the batch holds. Left, marked with a cross and titled sorted, shuffle=False: the dark cells run in a staircase down the diagonal, so each batch holds only one or two classes, noted in vermillion as classes in a batch: 1 or 2. Right, marked with a tick and titled shuffle=True: every cell is lightly shaded, noted in green as classes in a batch: all 10.
DL01.9.2-F5

Wrong beside right. If the training images arrive sorted by class, as files kept in one folder per class often do, unshuffled batches of 64 hold one or two classes each. Shuffled, every batch holds all ten.

“My data is already in a fixed order, so shuffle=False saves time.”

The batch gradient is an unbiased estimate of the full gradient only if every example is equally likely to be in the batch (ML01.5.4). A batch of only 3s pulls the network towards “everything is a 3”, the next batch of only 4s pulls it away, and the loss swings from batch to batch. Shuffle the training loader; the cost of a permutation of 1,077 numbers is nothing next to one step.

“drop_last=True is tidy, so I use it for validation too.”

On the 360 validation images with batches of 64 it keeps \({5 \cdot 64} = 320\) and never scores the other 40, a different 40 if the loader also shuffles. Validation and test loaders keep every example: drop_last=False and no shuffling, which changes nothing about the scores, since every example is visited anyway.

“I set torch.manual_seed, so my loader is repeatable.”

Without a generator argument the sampler draws its seed from PyTorch’s global random state, which everything else that uses random numbers also moves, so adding one line of code elsewhere changes the batches. Give the training loader its own seeded generator, as DL01.8.4 did.

Check yourself

Recall

1. How many steps does an epoch of \(n\) examples have with batches of \(\abs{B}\), without and with drop_last?

Show answer

\(\lceil n / \abs{B} \rceil\) without, \(\lfloor n / \abs{B} \rfloor\) with drop_last=True.

2. What does a seeded torch.Generator give a loader?

Show answer

A random source of its own, so the sequence of permutations, and with it every batch, is the same in every run, while each epoch still gets a new order.

3. What does the default collate function return for a list of \(\abs{B}\) pairs \({(\vect{x}_i, y_i)}\) of shapes (64,) and ()?

Show answer

A list of two tensors: the inputs stacked to shape (|B|, 64) and the labels to shape (|B|,).

4. With shuffling and drop_last=True, what is the chance that a given example sits out an epoch?

Show answer

\(r / n\), where \(r\) is the remainder of \(n\) divided by \(\abs{B}\): for 1,077 images in batches of 64, \(53 / 1077 \approx 0.049\).

5. Which DataLoader argument must not be set when you pass a sampler?

Show answer

shuffle: the sampler already fixes the order.

Exercises

1. ★ The 360 validation images go through a loader with batches of 64, and the 1,077 training images with batches of 50. Find the steps per epoch and the last batch size of each, with drop_last=False.

Show solution

Validation: \({360 = 5 \cdot 64 + 40}\), so 6 steps, the last of 40 images. Training: \({1077 = 21 \cdot 50 + 27}\), so 22 steps, the last of 27.

2. ★★ For which batch sizes from 1 to 100 does the training set of 1,077 images split into full batches only? What does that say about drop_last for those sizes?

Show solution

The remainder is zero exactly when \(\abs{B}\) divides 1,077. Since \({1077 = 3 \cdot 359}\) and 359 is prime, the divisors up to 100 are 1 and 3. For those two sizes drop_last changes nothing; for every other size from 2 to 100 a short last batch appears, and drop_last=True removes it.

3. ★★★ Training runs for \(E\) epochs with shuffling and drop_last=True. Show that the chance a given example is used in no epoch at all is \({(r/n)^E}\), and evaluate it for the 1,077 images in batches of 512 over 10 epochs.

Show solution

Each epoch draws a new permutation from the generator, and the example is left out of that epoch exactly when it lands in one of the last \(r\) of the \(n\) places, which happens with probability \(r/n\). The permutations of different epochs are drawn one after another from the generator and behave as independent, so the chance of being left out every time is the product, \({(r/n)^E}\). For \(\abs{B} = 512\), \(r = 53\) and \({(53/1077)^{10}} \approx 8.3 \times 10^{-14}\): in practice every example is seen, but in each single epoch about 5 percent are not.

Summary card

A loader shuffles the indices, cuts them into batches and stacks the items.

The anchor figure again: ten indices shuffled by randperm into 4 1 7 5 3 9 0 8 6 2, cut into batches of four, four and a short two, and stacked into tensors of shapes (4, 64), (4, 64) and (2, 64).
DL01.9.2-F1

Shuffle, cut, stack.

  • Steps per epoch. \(\lceil n / \abs{B} \rceil\), or \(\lfloor n / \abs{B} \rfloor\) with drop_last.
  • 1,077 images, \(\abs{B} = 64\). 17 steps, the last of 53; 16 with drop_last.
  • Repeatable. DataLoader(..., shuffle=True, generator=torch.Generator().manual_seed(0)).
  • Collate. A list of \(\abs{B}\) pairs becomes a list of two stacked tensors.

Sources and next steps

Sources. Each was opened and checked against this lesson on 15 November 2026.

  • PyTorch 2.15 documentation, torch.utils.data: batch_size defaults to 1 and shuffle to False, which reshuffles the data at every epoch when True; drop_last=True drops the last incomplete batch, and otherwise the last batch is smaller; the generator is used by RandomSampler to generate random indexes; a sampler excludes shuffle; default_collate prepends a batch dimension, converts NumPy arrays and Python numbers into tensors and preserves the data structure, returning a list for a batch of tuples; the BatchSampler example on range(10) with batch size 3.
  • S. J. D. Prince, Understanding Deep Learning, MIT Press 2023, section 6.2.1: batches are usually drawn without replacement, and a single pass through the whole training set is an epoch; figure 7.8 trains a network from a DataLoader(TensorDataset(x, y), batch_size=10, shuffle=True).
  • The behaviour described in the collapsed block was read in the source code of PyTorch 2.15 (torch/utils/data/dataloader.py and sampler.py) and checked by the cell that follows it.

Next. DL01.9.3 Batch size and the noise of a gradient asks what the batch size does to the gradient each step follows.

This lesson unlocks

DL01.9.3 Batch size and the noise of a gradient

Your progress is saved in this browser only. No account, no tracking.
Back to top