Batching and Training Loops in PyTorch

Training a model on a dataset rarely means handing the model every sample at once. Instead, the samples are fed in groups called batches, and each pass of the training loop works on one batch at a time. A data loader is the object that produces these batches from a dataset. In PyTorch, the DataLoader class from torch.utils.data does this job. The source for this lesson shows it being configured for three separate datasets: training, validation, and testing.

This article starts from a working configuration and asks what each argument does. You will form a prediction about the batches a loader produces, then check that prediction against the rules that govern the DataLoader class. The goal is not to memorize a template. The goal is to be able to read any DataLoader call and say what batches it will produce.

The DataLoader class is defined by a set of arguments. The general form given in the source is DataLoader(dataset, batch_size, shuffle, num_workers, drop_last). Each argument controls one aspect of the batches the loop receives. The dataset is the source of samples. The other four arguments are the levers: batch_size sets how many samples go into each batch, shuffle controls whether the samples are reordered, num_workers sets how many data-loading subprocesses run, and drop_last controls whether a smaller final batch is discarded.

Because each argument affects a different part of what the loop eventually receives, the parameters are best learned as a set of separate levers rather than as one combined rule. When a batch looks wrong, you do not need to change everything. You need to find the one lever whose effect matches the symptom.

The batch_size argument sets the number of samples in each batch. Think of the loader as walking through the dataset in order and cutting it into consecutive slices of that size. The diagram below shows a dataset of 10 samples, indexed 0 through 9, with batch_size set to 3.

Notice that the batch sizes are not always equal. The last batch can be smaller than batch_size whenever the dataset size is not an exact multiple of batch_size. This is the situation that drop_last is designed to handle.

The drop_last argument is the one most likely to surprise a learner, because it changes the amount of data a single pass uses without changing the dataset itself. When drop_last is True, a final batch smaller than batch_size is never yielded, so the samples in that batch are skipped for that pass. When it is False, the smaller batch is yielded like any other.

The shuffle argument controls whether data is shuffled. When shuffle is enabled, the samples are reordered before they are grouped into batches, so the same samples can land in different batches from one pass to the next. Without shuffling, the loader keeps the dataset's original order, and the same grouping appears each time.

The sample orders shown in the diagram are illustrative values chosen for this lesson, not output from a real run. What matters is the pattern: with shuffle enabled, consecutive passes can group the same samples differently, while the set of samples in the dataset stays the same.

The num_workers argument sets the number of data-loading subprocesses. These subprocesses are separate from the main process that runs the training loop. Their job is to prepare the data, so the loop has batches available to consume. The data flows from the dataset to the worker subprocesses, which produce batches, and then the training loop consumes those batches one at a time.

The source describes num_workers only as the number of data-loading subprocesses. It does not walk through the details of queueing. So when reasoning about a loader, treat num_workers as the setting that controls how many subprocesses load data, and keep it separate from batch_size and drop_last, which decide the content and count of the batches themselves.

Placing the three calls side by side makes the configuration choices easier to see. Most arguments are shared, and the differences are deliberate. The table below lists what each loader receives from the source listing.

The pattern is that only the training loader is shuffled and only the training loader drops its final partial batch. The source shows these settings without giving the reasoning behind them, so treat the table as a record of this configuration rather than a rule for every project. For the validation and test loaders, shuffle is simply not passed, so their order follows whatever default the DataLoader class uses when that argument is omitted. The source does not state that default, so check the DataLoader documentation before relying on it.

Changing batch_size changes how many batches a pass contains. Because each pass of the training loop works on one batch at a time, the number of batches per pass also determines how many times the loop steps through the data. The table below compares a dataset of 12 samples at two batch sizes, with drop_last set to False.

The source does not discuss how batch size affects training outcomes, so this article limits itself to the counting consequence. If the loop takes one step per batch, then doubling batch_size roughly halves the number of steps per pass, and each step handles twice the samples.

When batches are unexpected, the source suggests checking each argument against its effect, starting from the one whose behavior matches the symptom. A wrong sample order points to the shuffle step. A batch with an unexpected number of samples points to batch_size. A missing final batch points to drop_last. Use the lookup below as a quick trace.

A common surprise is a batch count that is one lower than expected. Before suspecting the dataset, check whether drop_last is set on that loader. The training loader in the source is the one with drop_last set, so a shortfall there has a different cause than the same shortfall in the validation or test loader.