Writing a training loop from scratch in PyTorch

Author: fchollet Date created: 2023/06/25 Last modified: 2023/06/25 Description: Writing low-level training & evaluation loops in PyTorch.
View in Colab • GitHub source
Keras provides default training and evaluation loops, fit() and evaluate(). Their usage is covered in the guide Training & evaluation with the built-in methods.
If you want to customize the learning algorithm of your model while still leveraging the convenience of fit() (for instance, to train a GAN using fit()), you can subclass the Model class and implement your own train_step() method, which is called repeatedly during fit().
Now, if you want very low-level control over training & evaluation, you should write your own training & evaluation loops from scratch. This is what this guide is about.
To write a custom training loop, we need the following ingredients:
Let's line them up. We'll use torch-native objects in each case – except, of course, for the Keras model.
First, let's get the model and the MNIST dataset:
Next, here's our PyTorch optimizer and our PyTorch loss function:
Let's train our model using mini-batch gradient with a custom training loop.
Calling loss.backward() on a loss tensor triggers backpropagation. Once that's done, your optimizer is magically aware of the gradients for each variable and can update its variables, which is done via optimizer.step(). Tensors, variables, optimizers are all interconnected to one another via hidden global state. Also, don't forget to call model.zero_grad() before loss.backward(), or you won't get the right gradients for your variables.
Here's our training loop, step by step:
As an alternative, let's look at what the loop looks like when using a Keras optimizer and a Keras loss function.
Important differences:
Also, a big gotcha: while all NumPy/TensorFlow/JAX/Keras APIs as well as Python unittest APIs use the argument order convention fn(y_true, y_pred) (reference values first, predicted values second), PyTorch actually uses fn(y_pred, y_true) for its losses. So make sure to invert the order of logits and targets.
Let's add metrics monitoring to this basic training loop.
You can readily reuse built-in Keras metrics (or custom ones you wrote) in such training loops written from scratch. Here's the flow:
Let's use this knowledge to compute CategoricalAccuracy on training and validation data at the end of each epoch:
Here's our training & evaluation loop:
Layers & models recursively track any losses created during the forward pass by layers that call self.add_loss(value). The resulting list of scalar loss values are available via the property model.losses at the end of the forward pass.
If you want to be using these loss components, you should sum them and add them to the main loss in your training step.
Consider this layer, that creates an activity regularization loss:
Let's build a really simple model that uses it:
Here's what our training loop should look like now: