title writing a training loop from scratch in jax fchollethttps twitter comfchollet date created 20230625 last modified 20230625 description writing lowlevel training evaluation loops in jax accelerator none setup this guide can only be run with the jax backend we import tf so we can use tf data introduction keras provides default training and evaluation loops fit and evaluate their usage is covered in the guide training evaluation with the builtin methodshttps keras ioguidestrainingwithbuiltinmethods 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 trainstep method which is called repeatedly during fit now if you want very lowlevel control over training evaluation you should write your own training evaluation loops from scratch this is what this guide is about a first endtoend example to write a custom training loop we need the following ingredients a model to train of course an optimizer you could either use an optimizer from keras optimizers or one from the optax package a loss function a dataset the standard in the jax ecosystem is to load data via tf data so that s what we ll use let s line them up first let s get the model and the mnist dataset prepare the training dataset reserve 10 000 samples for validation prepare the training dataset prepare the validation dataset next here s the loss function and the optimizer we ll use a keras optimizer in this case instantiate a loss function instantiate an optimizer getting gradients in jax let s train our model using minibatch gradient with a custom training loop in jax gradients are computed via metaprogramming you call the jax grad or jax valueandgrad on a function in order to create a gradientcomputing 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 python def computelossx y return loss once you have such a function you can compute gradients via metaprogramming as such python gradfn jax gradcomputeloss grads gradfnx y 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 valueandgrad instead of jax grad python gradfn jax valueandgradcomputeloss loss grads gradfnx y jax computation is purely stateless 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 nontrainable 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 python def computelossandupdatestrainablevariables nontrainablevariables x y return loss nontrainablevariables once you have such a function you can get the gradient function by specifying haxaux in valueandgrad 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 python gradfn jax valueandgradcomputelossandupdates hasauxtrue loss nontrainablevariables grads gradfn trainablevariables nontrainablevariables x y now that we have established the basics let s implement this computelossandupdates function keras models have a statelesscall 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 nontrainable variables let s get the gradient function the training step function next let s implement the endtoend 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 trainablevariables and nontrainablevariables the model s variables optimizervariables the optimizer s state variables such as momentum accumulators to update the trainable variables we use the optimizer s stateless method statelessapply it s equivalent to optimizer apply but it requires always passing trainablevariables and optimizervariables it returns both the updated trainable variables and the updated optimizervariables return updated state make it fast with jax jit 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 trainstep 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 functionjitcompiletrue let s try it return updated state we re now ready to train our model the training loop itself is trivial we just repeatedly call loss state trainstepstate data note we convert the tf tensors yielded by the tf data dataset to numpy before passing them to our jax function all variables must be built beforehand the model must be built and the optimizer must be built since we re using a functional api model it s already built but if it were a subclassed model you d need to call it on a batch of data to build it build optimizer variables training loop log every 100 batches 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 assignnewvalue on each model variable you want to update lowlevel handling of metrics let s add metrics monitoring to this basic training loop you can readily reuse builtin keras metrics or custom ones you wrote in such training loops written from scratch here s the flow instantiate the metric at the start of the loop include metricvariables in the trainstep arguments and computelossandupdates arguments call metric statelessupdatestate in the computelossandupdates function it s equivalent to updatestate only stateless when you need to display the current value of the metric outside the trainstep in the eager scope attach the new metric variable values to the metric object and vall metric result call metric resetstate when you need to clear the state of the metric typically at the end of an epoch let s use this knowledge to compute categoricalaccuracy on training and validation data at the end of training get a fresh model instantiate an optimizer to train the model instantiate a loss function prepare the metrics return updated state we ll also prepare an evaluation step function here are our loops build optimizer variables training loop log every 100 batches eval loop log every 100 batches lowlevel handling of losses tracked by the model layers models recursively track any losses created during the forward pass by layers that call self addlossvalue 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 insert activity regularization as a layer here s what our computelossandupdates function should look like now pass returnlossestrue to model statelesscall sum the resulting losses and add them to the main loss that s it setup this guide can only be run with the jax backend we import tf so we can use tf data introduction keras provides default training and evaluation loops fit and evaluate their usage is covered in the guide training evaluation with the built in methods https keras io guides training_with_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 a first end to end example to write a custom training loop we need the following ingredients a model to train of course an optimizer you could either use an optimizer from keras optimizers or one from the optax package a loss function a dataset the standard in the jax ecosystem is to load data via tf data so that s what we ll use let s line them up first let s get the model and the mnist dataset prepare the training dataset reserve 10 000 samples for validation prepare the training dataset prepare the validation dataset next here s the loss function and the optimizer we ll use a keras optimizer in this case instantiate a loss function instantiate an optimizer getting gradients in jax 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 python def compute_loss x y return loss once you have such a function you can compute gradients via metaprogramming as such python grad_fn jax grad compute_loss grads grad_fn x y 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 python grad_fn jax value_and_grad compute_loss loss grads grad_fn x y jax computation is purely stateless 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 python def compute_loss_and_updates trainable_variables non_trainable_variables x y return loss non_trainable_variables once you have such a function you can get the gradient function by specifying hax_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 python grad_fn jax value_and_grad compute_loss_and_updates has_aux true loss non_trainable_variables grads grad_fn trainable_variables non_trainable_variables x y 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 the training step 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 trainable_variables and non_trainable_variables the model s variables optimizer_variables the optimizer s state variables such as momentum accumulators 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 return updated state make it fast with jax jit 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 return updated state we re now ready to train our model the training loop itself is trivial we just repeatedly call loss state train_step state data note we convert the tf tensors yielded by the tf data dataset to numpy before passing them to our jax function all variables must be built beforehand the model must be built and the optimizer must be built since we re using a functional api model it s already built but if it were a subclassed model you d need to call it on a batch of data to build it build optimizer variables training loop log every 100 batches 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 low level handling of metrics 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 instantiate the metric at the start of the loop include metric_variables in the train_step arguments and compute_loss_and_updates arguments call metric stateless_update_state in the compute_loss_and_updates function it s equivalent to update_state only stateless when you need to display the current value of the metric outside the train_step in the eager scope attach the new metric variable values to the metric object and vall metric result call metric reset_state when you need to clear the state of the metric typically at the end of an epoch let s use this knowledge to compute categoricalaccuracy on training and validation data at the end of training get a fresh model instantiate an optimizer to train the model instantiate a loss function prepare the metrics return updated state we ll also prepare an evaluation step function here are our loops build optimizer variables training loop log every 100 batches eval loop log every 100 batches low level handling of losses tracked by the model 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 insert activity regularization as a layer here s what our compute_loss_and_updates function should look like now pass return_losses true to model stateless_call sum the resulting losses and add them to the main loss that s it
import os


os.environ["KERAS_BACKEND"] = "jax"

import jax


import tensorflow as tf
import keras
import numpy as np






def get_model():
    inputs = keras.Input(shape=(784,), name="digits")
    x1 = keras.layers.Dense(64, activation="relu")(inputs)
    x2 = keras.layers.Dense(64, activation="relu")(x1)
    outputs = keras.layers.Dense(10, name="predictions")(x2)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


model = get_model()


batch_size = 32
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = np.reshape(x_train, (-1, 784)).astype("float32")
x_test = np.reshape(x_test, (-1, 784)).astype("float32")
y_train = keras.utils.to_categorical(y_train)
y_test = keras.utils.to_categorical(y_test)


x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]


train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)


val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(batch_size)




loss_fn = keras.losses.CategoricalCrossentropy(from_logits=True)


optimizer = keras.optimizers.Adam(learning_rate=1e-3)




def compute_loss_and_updates(
    trainable_variables, non_trainable_variables, x, y
):
    y_pred, non_trainable_variables = model.stateless_call(
        trainable_variables, non_trainable_variables, x
    )
    loss = loss_fn(y, y_pred)
    return loss, non_trainable_variables




grad_fn = jax.value_and_grad(compute_loss_and_updates, has_aux=True)




def train_step(state, data):
    trainable_variables, non_trainable_variables, optimizer_variables = state
    x, y = data
    (loss, non_trainable_variables), grads = grad_fn(
        trainable_variables, non_trainable_variables, x, y
    )
    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        optimizer_variables, grads, trainable_variables
    )
    
    return loss, (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    )





@jax.jit
def train_step(state, data):
    trainable_variables, non_trainable_variables, optimizer_variables = state
    x, y = data
    (loss, non_trainable_variables), grads = grad_fn(
        trainable_variables, non_trainable_variables, x, y
    )
    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        optimizer_variables, grads, trainable_variables
    )
    
    return loss, (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    )





optimizer.build(model.trainable_variables)

trainable_variables = model.trainable_variables
non_trainable_variables = model.non_trainable_variables
optimizer_variables = optimizer.variables
state = trainable_variables, non_trainable_variables, optimizer_variables


for step, data in enumerate(train_dataset):
    data = (data[0].numpy(), data[1].numpy())
    loss, state = train_step(state, data)
    
    if step % 100 == 0:
        print(f"Training loss (for 1 batch) at step {step}: {float(loss):.4f}")
        print(f"Seen so far: {(step + 1) * batch_size} samples")



trainable_variables, non_trainable_variables, optimizer_variables = state
for variable, value in zip(model.trainable_variables, trainable_variables):
    variable.assign(value)
for variable, value in zip(
    model.non_trainable_variables, non_trainable_variables
):
    variable.assign(value)




model = get_model()


optimizer = keras.optimizers.Adam(learning_rate=1e-3)

loss_fn = keras.losses.CategoricalCrossentropy(from_logits=True)


train_acc_metric = keras.metrics.CategoricalAccuracy()
val_acc_metric = keras.metrics.CategoricalAccuracy()


def compute_loss_and_updates(
    trainable_variables, non_trainable_variables, metric_variables, x, y
):
    y_pred, non_trainable_variables = model.stateless_call(
        trainable_variables, non_trainable_variables, x
    )
    loss = loss_fn(y, y_pred)
    metric_variables = train_acc_metric.stateless_update_state(
        metric_variables, y, y_pred
    )
    return loss, (non_trainable_variables, metric_variables)


grad_fn = jax.value_and_grad(compute_loss_and_updates, has_aux=True)


@jax.jit
def train_step(state, data):
    (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
        metric_variables,
    ) = state
    x, y = data
    (loss, (non_trainable_variables, metric_variables)), grads = grad_fn(
        trainable_variables, non_trainable_variables, metric_variables, x, y
    )
    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        optimizer_variables, grads, trainable_variables
    )
    
    return loss, (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
        metric_variables,
    )





@jax.jit
def eval_step(state, data):
    trainable_variables, non_trainable_variables, metric_variables = state
    x, y = data
    y_pred, non_trainable_variables = model.stateless_call(
        trainable_variables, non_trainable_variables, x
    )
    loss = loss_fn(y, y_pred)
    metric_variables = val_acc_metric.stateless_update_state(
        metric_variables, y, y_pred
    )
    return loss, (
        trainable_variables,
        non_trainable_variables,
        metric_variables,
    )





optimizer.build(model.trainable_variables)

trainable_variables = model.trainable_variables
non_trainable_variables = model.non_trainable_variables
optimizer_variables = optimizer.variables
metric_variables = train_acc_metric.variables
state = (
    trainable_variables,
    non_trainable_variables,
    optimizer_variables,
    metric_variables,
)


for step, data in enumerate(train_dataset):
    data = (data[0].numpy(), data[1].numpy())
    loss, state = train_step(state, data)
    
    if step % 100 == 0:
        print(f"Training loss (for 1 batch) at step {step}: {float(loss):.4f}")
        _, _, _, metric_variables = state
        for variable, value in zip(
            train_acc_metric.variables, metric_variables
        ):
            variable.assign(value)
        print(f"Training accuracy: {train_acc_metric.result()}")
        print(f"Seen so far: {(step + 1) * batch_size} samples")

metric_variables = val_acc_metric.variables
(
    trainable_variables,
    non_trainable_variables,
    optimizer_variables,
    metric_variables,
) = state
state = trainable_variables, non_trainable_variables, metric_variables


for step, data in enumerate(val_dataset):
    data = (data[0].numpy(), data[1].numpy())
    loss, state = eval_step(state, data)
    
    if step % 100 == 0:
        print(
            f"Validation loss (for 1 batch) at step {step}: {float(loss):.4f}"
        )
        _, _, metric_variables = state
        for variable, value in zip(val_acc_metric.variables, metric_variables):
            variable.assign(value)
        print(f"Validation accuracy: {val_acc_metric.result()}")
        print(f"Seen so far: {(step + 1) * batch_size} samples")




class ActivityRegularizationLayer(keras.layers.Layer):
    def call(self, inputs):
        self.add_loss(1e-2 * jax.numpy.sum(inputs))
        return inputs




inputs = keras.Input(shape=(784,), name="digits")
x = keras.layers.Dense(64, activation="relu")(inputs)

x = ActivityRegularizationLayer()(x)
x = keras.layers.Dense(64, activation="relu")(x)
outputs = keras.layers.Dense(10, name="predictions")(x)

model = keras.Model(inputs=inputs, outputs=outputs)




def compute_loss_and_updates(
    trainable_variables, non_trainable_variables, metric_variables, x, y
):
    y_pred, non_trainable_variables, losses = model.stateless_call(
        trainable_variables, non_trainable_variables, x, return_losses=True
    )
    loss = loss_fn(y, y_pred)
    if losses:
        loss += jax.numpy.sum(losses)
    metric_variables = train_acc_metric.stateless_update_state(
        metric_variables, y, y_pred
    )
    return loss, non_trainable_variables, metric_variables