Writing a training loop from scratch in JAX

Author: fchollet Date created: 2023/06/25 Last modified: 2023/06/25 Description: Writing low-level training & evaluation loops in JAX.
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:
First, let's get the model and the MNIST dataset:
Next, here's the loss function and the optimizer. We'll use a Keras optimizer in this case.
Let's train our model using mini-batch gradient with a custom training loop.
In JAX, gradients are computed via metaprogramming: you call the jax.grad (or jax.value_and_grad on a function in order to create a gradient-computing function for that first function.
So the first thing we need is a function that returns the loss value. That's the function we'll use to generate the gradient function. Something like this:
Once you have such a function, you can compute gradients via metaprogramming as such:
Typically, you don't just want to get the gradient values, you also want to get the loss value. You can do this by using jax.value_and_grad instead of jax.grad:
In JAX, everything must be a stateless function – so our loss computation function must be stateless as well. That means that all Keras variables (e.g. weight tensors) must be passed as function inputs, and any variable that has been updated during the forward pass must be returned as function output. The function have no side effect.
During the forward pass, the non-trainable variables of a Keras model might get updated. These variables could be, for instance, RNG seed state variables or BatchNormalization statistics. We're going to need to return those. So we need something like this:
Once you have such a function, you can get the gradient function by specifying has_aux in value_and_grad: it tells JAX that the loss computation function returns more outputs than just the loss. Note that the loss should always be the first output.
Now that we have established the basics, let's implement this compute_loss_and_updates function. Keras models have a stateless_call method which will come in handy here. It works just like model.__call__, but it requires you to explicitly pass the value of all the variables in the model, and it returns not just the __call__ outputs but also the (potentially updated) non-trainable variables.
Let's get the gradient function:
Next, let's implement the end-to-end training step, the function that will both run the forward pass, compute the loss, compute the gradients, but also use the optimizer to update the trainable variables. This function also needs to be stateless, so it will get as input a state tuple that includes every state element we're going to use:
To update the trainable variables, we use the optimizer's stateless method stateless_apply. It's equivalent to optimizer.apply(), but it requires always passing trainable_variables and optimizer_variables. It returns both the updated trainable variables and the updated optimizer_variables.
By default, JAX operations run eagerly, just like in TensorFlow eager mode and PyTorch eager mode. And just like TensorFlow eager mode and PyTorch eager mode, it's pretty slow – eager mode is better used as a debugging environment, not as a way to do any actual work. So let's make our train_step fast by compiling it.
When you have a stateless JAX function, you can compile it to XLA via the @jax.jit decorator. It will get traced during its first execution, and in subsequent executions you will be executing the traced graph (this is just like @tf.function(jit_compile=True). Let's try it:
We're now ready to train our model. The training loop itself is trivial: we just repeatedly call loss, state = train_step(state, data).
A key thing to notice here is that the loop is entirely stateless – the variables attached to the model (model.weights) are never getting updated during the loop. Their new values are only stored in the state tuple. That means that at some point, before saving the model, you should be attaching the new variable values back to the model.
Just call variable.assign(new_value) on each model variable you want to update:
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 training:
We'll also prepare an evaluation step function:
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 compute_loss_and_updates function should look like now: