A batch, or mini-batch, is a small group of samples processed together. A model does not need to process the entire training dataset at once. In a batch tensor, the first axis counts samples; it is called the batch axis. section 2.2.7
Batching handwritten digits
With a batch size of 128, `train_images@r[1:128, , ]` selects one group of MNIST digit images; `train_images@r[129:256, , ]` selects the next. During training, each group is paired with its corresponding labels. A training loop can then call `fit(model, train_images, train_labels, epochs = 10, batch_size = 128)`. section 2.2.7 section 2.5.3
One mini-batch stochastic gradient descent step does the following:
- Draw samples and their targets, then run the model to get predictions—the forward pass. section 2.4.3
- Measure prediction error with a loss, then compute the loss gradient with respect to the model parameters—the backward pass. section 2.4.3
- Move the parameters a little in the direction opposite the gradient. The learning rate controls the size of this move. section 2.4.3
An epoch repeats that training step for each batch in the training data. Repeating epochs makes the full training loop. Batches are commonly small—often 8 to 128 samples—but the forecasting example below uses 256, so 128 is not a required size. section 2.5.3 section 4.0 section 13.2.1
fit <- function(model, images, labels, epochs, batch_size = 128) {
for (epoch_counter in seq_len(epochs)) {
cat("Epoch ", epoch_counter, "\n")
batch_generator <- new_batch_generator(images, labels, batch_size)
batch_counter <- 0
repeat {
batch <- batch_generator() %||% break
batch_counter <- batch_counter + 1
loss <- one_training_step(model, batch$images, batch$labels)
if (batch_counter %% 100 == 0)
cat(sprintf("loss at batch %s: %.2f\n", batch_counter, loss))
}
}
}
This R training loop creates a batch generator for each epoch, takes batches until none remain, and calls `one_training_step()` on each batch. section 2.5.3
Batching also applies when samples are generated rather than stored individually. For temperature forecasting, `timeseries_dataset_from_array()` extracts overlapping input windows and their future-temperature targets on the fly. Its `batch_size` argument groups those windows; one inspected batch has sample shape `(256, 120, 14)` and target shape `(256)`. section 13.2.1
