To train the image classification model and select optimal hyperparameters based on validation set performance, a comprehensive training function is defined. This function orchestrates the training loop over a specified number of epochs using data parallelism across multiple devices. It initializes an optimization algorithm, such as stochastic gradient descent (SGD) with momentum and weight decay, and incorporates a step learning rate scheduler to periodically decay the learning rate by a specified factor. Within each epoch, the function iterates through training mini-batches to update parameters and accumulate training loss and accuracy. If a validation iterator is provided, it evaluates the model's accuracy on the hold-out validation set, plotting these metrics dynamically to monitor the model's generalization progress.

Claude

To evaluate the predictions of the ResNet-18 model on the CIFAR-10 dataset, the cross-entropy loss function is utilized. In some framework implementations, this loss function is explicitly configured to compute the unreduced loss (e.g., `reduction="none"`), returning individual loss values for each example rather than a single aggregated scalar.

CIFAR-10 Cross-Entropy Loss Configuration

While the training loop is responsible for optimizing a model's parameters, a validation dataset is crucial for measuring the model's quality on unseen data. A standard practice is to pass a validation dataloader through the model once during each epoch. By evaluating the model's performance without updating its parameters—often in a no-gradient context—practitioners can monitor generalization and detect overfitting during the iterative training process.

Validation Dataset Model Evaluation

In a data parallelism setup, the training process follows a specific sequence during each iteration. First, a random minibatch of training data is split into $$k$$ equal portions and distributed across the $$k$$ available GPUs. Each GPU then independently calculates the loss and the local gradients of the model parameters using its assigned data subset. These local gradients from all $$k$$ GPUs are subsequently aggregated (via an allreduce operation) to compute the overall minibatch stochastic gradient. This aggregated gradient is redistributed to every GPU, and each GPU uses it to independently update its complete set of model parameters, ensuring synchronization across the system.

Data Parallelism Training Process

The main training loop for a deep learning model executes a systematic, iterative process to optimize parameters $$\mathbf{w}$$ and $$b$$. During each epoch, the loop passes through the entire training dataset. For every iteration within an epoch, a minibatch $$\mathcal{B}$$ is processed. The model computes the loss for this minibatch, averages it over the examples in the batch, and calculates the gradients of the averaged loss with respect to each parameter using $$\mathbf{g} \leftarrow \partial_{(\mathbf{w},b)} \frac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} l(\mathbf{x}^{(i)}, y^{(i)}, \mathbf{w}, b)$$. Because the loss is averaged per minibatch, the gradient in the optimization algorithm does not need to be separately divided by the batch size. Finally, the optimization algorithm updates the parameters using the rule $$(\mathbf{w}, b) \leftarrow (\mathbf{w}, b) - \eta \mathbf{g}$$. This cycle repeats until the training is complete.

Deep Learning Minibatch Training Loop

A piecewise constant learning rate schedule, often referred to as a multi-factor scheduler, maintains a steady learning rate for predefined intervals of training and abruptly drops the rate at specific milestones. Mathematically, given a set of milestone times $$s$$ (such as $$s = \{5, 10, 20\}$$), the learning rate is updated according to the rule $$\eta_{t+1} \leftarrow \eta_t \cdot \alpha$$ whenever the current step $$t \in s$$, where $$\alpha$$ is the designated decay factor.

Piecewise Constant Learning Rate Schedule

Dive into Deep Learning

CIFAR-10 Model Training Function Definition

The `train_batch` function implements data-parallel training for a single minibatch across multiple GPUs. The procedure follows four sequential stages:

1. **Data splitting:** The minibatch of features and labels is divided across the available devices using `split_batch`.
2. **Per-GPU forward pass and loss:** Each GPU independently computes the model's output and the loss on its local data shard. The losses are summed per device.
3. **Per-GPU backpropagation:** Backpropagation is performed separately on each GPU to compute local gradients.
4. **Gradient synchronization and update:** An `allreduce` operation sums and broadcasts all gradients across GPUs within a `torch.no_grad()` context. Finally, each GPU independently updates its own copy of the model parameters using SGD, scaling the update by the full (unsplit) batch size.

```python
def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:
        l.backward()
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce([device_params[c][i].grad
                       for c in range(len(devices))])
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0])
```

Because there are no cross-device dependencies within the computational graph for a single minibatch, the per-GPU computations execute in parallel automatically.

Multi-GPU Minibatch Training Implementation

Deep learning frameworks can automatically parallelize independent computations across multiple GPUs without requiring explicit multi-threading or scheduling code from the user. This automatic parallelism is a direct consequence of asynchronous execution: when the frontend issues operations targeting different GPU devices sequentially, these operations are placed into separate backend queues for each device. Because no data dependency exists between operations on different devices, the backend processes them concurrently. However, if a synchronization barrier (such as torch.cuda.synchronize() or npx.waitall()) is inserted between the operations on the two devices, it forces the first device's work to complete before the second device's work begins, serializing execution and preventing parallelism.

Automatic Multi-GPU Parallelism via Asynchronous Execution

Distributed parallel training across multiple machines involves a synchronized sequence of operations to manage parameters over a comparatively lower bandwidth network fabric. The process follows seven steps: First, a batch of data is read on each machine, split, and transferred to local GPUs where predictions and gradients are computed separately. Second, the local gradients are aggregated (or partially aggregated) on the GPUs within each machine. Third, these aggregated gradients are sent from the GPUs to the machine's CPUs. Fourth, the CPUs transmit the gradients over the network to a central parameter server, which aggregates all gradients across the different machines. Fifth, the central server uses the aggregate gradients to update the parameters and broadcasts them back to the individual CPUs. Sixth, the CPUs send the updated parameters to one or more local GPUs. Finally, the parameters are spread across all GPUs on the machine for the next iteration.

Multi-Machine Distributed Parallel Training Process

When training a model such as linear regression from scratch on a synthetic dataset, the true underlying parameters are perfectly known. This provides a unique opportunity to evaluate the success of the training loop by directly comparing the ground truth weights and bias with the parameters learned by the model. While exactly recovering the true parameters is generally difficult or impossible for deep networks because unique solutions often do not exist, stochastic gradient descent typically discovers parameter configurations that yield highly accurate predictions.

Linear Regression Synthetic Data Training Evaluation

The `fit_epoch` method, typically registered within an object-oriented `Trainer` class, encapsulates the computations required for a single full pass over the training dataset. It begins by setting the model to training mode and iterating through the training dataloader to compute the loss via the `training_step` method. Within this loop, it zeroes the gradients, performs backpropagation, and updates the parameters. If a validation dataloader is provided, the method subsequently sets the model to evaluation mode and processes the validation batches within a no-gradient context to assess the model's performance.

Object-Oriented Fit Epoch Implementation

To execute the training loop for a model from scratch, we instantiate the required components and begin the optimization process. For example, we initialize a linear regression model, `LinearRegressionScratch`, specifying hyperparameters such as the learning rate, `lr=0.03`. We then create the dataset using a class like `SyntheticRegressionData`, passing in ground truth parameter values. Next, an object-oriented `Trainer` is instantiated with settings such as `max_epochs=3`. Finally, training is executed by calling `trainer.fit(model, data)`, which orchestrates the iterations over the dataset to optimize the model's parameters.

Linear Regression Training Execution Example

The train_ch11 function provides a generic, from-scratch training harness for benchmarking different optimization algorithms on a linear regression model. It accepts five arguments: a trainer function (trainer_fn), optimizer states, a hyperparameters dictionary, a data iterator, and the feature dimensionality. The function initializes weight parameters from a normal distribution with mean $$0$$ and standard deviation $$0.01$$, and a bias of $$0$$. During each epoch, it iterates over minibatches, computes the mean squared loss, performs backpropagation, and calls trainer_fn to update the parameters. It periodically evaluates the full-dataset loss and records the cumulative training time. This reusable design allows any optimizer with a compatible call signature to be plugged in by simply passing a different trainer_fn. A companion function, train_concise_ch11, provides an equivalent harness that delegates to framework-level optimizer APIs instead of custom update functions.

Generic Optimization Training Function

The fundamental intuition behind a piecewise constant learning rate schedule is to allow the optimization process to proceed with a steady learning rate until it reaches a stationary point or plateau regarding the distribution of weight vectors. By abruptly dropping the learning rate only after efficiently converging to this suitable solution, the model reduces the inherent variance of its parameters. This approach enables the discovery of a higher quality proxy to a good local minimum and slightly better overall solutions.

Intuition Behind Piecewise Constant Learning Rate Schedules

A multi-factor (or piecewise constant) learning rate scheduler decreases the learning rate by a specific factor at predefined milestone epochs. While modern frameworks offer built-in solutions—such as `torch.optim.lr_scheduler.MultiStepLR` in PyTorch and `lr_scheduler.MultiFactorScheduler` in MXNet—this strategy can also be implemented from scratch. The following Python code demonstrates a custom scheduler that multiplies the base learning rate by a given factor whenever the current epoch is found in the predefined set of milestone steps:

```python
class MultiFactorScheduler:
    def __init__(self, step, factor, base_lr):
        self.step = step
        self.factor = factor
        self.base_lr = base_lr

    def __call__(self, epoch):
        if epoch in self.step:
            self.base_lr = self.base_lr * self.factor
            return self.base_lr
        else:
            return self.base_lr
```

Multi-Factor Learning Rate Scheduler Implementation

To stabilize optimization and achieve better convergence during fine-tuning, a step learning rate scheduler is often employed to periodically reduce the learning rate. By defining a period (lr_period) and a decay factor (lr_decay), tools like torch.optim.lr_scheduler.StepLR automatically multiply the learning rate by the decay factor after every specified number of epochs, ensuring smaller optimization steps as the model approaches a local minimum on the target dataset.

Learn Before

Related