title customizing what happens in fit with jax fchollethttps twitter comfchollet date created 20230627 last modified 20230627 description overriding the training step of the model class with jax accelerator gpu introduction when you re doing supervised learning you can use fit and everything works smoothly when you need to take control of every little detail you can write your own training loop entirely from scratch but what if you need a custom training algorithm but you still want to benefit from the convenient features of fit such as callbacks builtin distribution support or step fusing a core principle of keras is progressive disclosure of complexity you should always be able to get into lowerlevel workflows in a gradual way you shouldn t fall off a cliff if the highlevel functionality doesn t exactly match your use case you should be able to gain more control over the small details while retaining a commensurate amount of highlevel convenience when you need to customize what fit does you should override the training step function of the model class this is the function that is called by fit for every batch of data you will then be able to call fit as usual and it will be running your own learning algorithm note that this pattern does not prevent you from building models with the functional api you can do this whether you re building sequential models functional api models or subclassed models let s see how that works setup this guide can only be run with the jax backend a first simple example let s start from a simple example we create a new class that subclasses keras model we implement a fullystateless computelossandupdates method to compute the loss as well as the updated values for the nontrainable variables of the model internally it calls statelesscall and the builtin computeloss we implement a fullystateless trainstep method to compute current metric values including the loss as well as updated values for the trainable variables the optimizer variables and the metric variables note that you can also take into account the sampleweight argument by unpacking the data as x y sampleweight data passing sampleweight to computeloss passing sampleweight alongside y and ypred to metrics in statelessupdatestate get the gradient function compute the gradients update trainable variables and optimizer variables update metrics return metric logs and updated state variables let s try this out construct and compile an instance of custommodel just use fit as usual going lowerlevel naturally you could just skip passing a loss function in compile and instead do everything manually in trainstep likewise for metrics here s a lowerlevel example that only uses compile to configure the optimizer get the gradient function compute the gradients update trainable variables and optimizer variables update metrics return metric logs and updated state variables we list our metric objects here so that resetstates can be called automatically at the start of each epoch or at the start of evaluate construct an instance of custommodel we don t passs a loss or metrics here just use fit as usual you can use callbacks etc providing your own evaluation step what if you want to do the same for calls to model evaluate then you would override teststep in exactly the same way here s what it looks like unpack the data compute predictions and loss update metrics return metric logs and updated state variables construct an instance of custommodel evaluate with our custom teststep that s it introduction when you re doing supervised learning you can use fit and everything works smoothly when you need to take control of every little detail you can write your own training loop entirely from scratch but what if you need a custom training algorithm but you still want to benefit from the convenient features of fit such as callbacks built in distribution support or step fusing a core principle of keras is progressive disclosure of complexity you should always be able to get into lower level workflows in a gradual way you shouldn t fall off a cliff if the high level functionality doesn t exactly match your use case you should be able to gain more control over the small details while retaining a commensurate amount of high level convenience when you need to customize what fit does you should override the training step function of the model class this is the function that is called by fit for every batch of data you will then be able to call fit as usual and it will be running your own learning algorithm note that this pattern does not prevent you from building models with the functional api you can do this whether you re building sequential models functional api models or subclassed models let s see how that works setup this guide can only be run with the jax backend a first simple example let s start from a simple example we create a new class that subclasses keras model we implement a fully stateless compute_loss_and_updates method to compute the loss as well as the updated values for the non trainable variables of the model internally it calls stateless_call and the built in compute_loss we implement a fully stateless train_step method to compute current metric values including the loss as well as updated values for the trainable variables the optimizer variables and the metric variables note that you can also take into account the sample_weight argument by unpacking the data as x y sample_weight data passing sample_weight to compute_loss passing sample_weight alongside y and y_pred to metrics in stateless_update_state get the gradient function compute the gradients update trainable variables and optimizer variables update metrics return metric logs and updated state variables let s try this out construct and compile an instance of custommodel just use fit as usual going lower level naturally you could just skip passing a loss function in compile and instead do everything manually in train_step likewise for metrics here s a lower level example that only uses compile to configure the optimizer get the gradient function compute the gradients update trainable variables and optimizer variables update metrics return metric logs and updated state variables we list our metric objects here so that reset_states can be called automatically at the start of each epoch or at the start of evaluate construct an instance of custommodel we don t passs a loss or metrics here just use fit as usual you can use callbacks etc providing your own evaluation step what if you want to do the same for calls to model evaluate then you would override test_step in exactly the same way here s what it looks like unpack the data compute predictions and loss update metrics return metric logs and updated state variables construct an instance of custommodel evaluate with our custom test_step that s it
import os


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

import jax
import keras
import numpy as np




class CustomModel(keras.Model):
    def compute_loss_and_updates(
        self,
        trainable_variables,
        non_trainable_variables,
        x,
        y,
        training=False,
    ):
        y_pred, non_trainable_variables = self.stateless_call(
            trainable_variables,
            non_trainable_variables,
            x,
            training=training,
        )
        loss = self.compute_loss(x, y, y_pred)
        return loss, (y_pred, non_trainable_variables)

    def train_step(self, state, data):
        (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        ) = state
        x, y = data

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

        
        (loss, (y_pred, non_trainable_variables)), grads = grad_fn(
            trainable_variables,
            non_trainable_variables,
            x,
            y,
            training=True,
        )

        
        (
            trainable_variables,
            optimizer_variables,
        ) = self.optimizer.stateless_apply(
            optimizer_variables, grads, trainable_variables
        )

        
        new_metrics_vars = []
        for metric in self.metrics:
            this_metric_vars = metrics_variables[
                len(new_metrics_vars) : len(new_metrics_vars)
                + len(metric.variables)
            ]
            if metric.name == "loss":
                this_metric_vars = metric.stateless_update_state(
                    this_metric_vars, loss
                )
            else:
                this_metric_vars = metric.stateless_update_state(
                    this_metric_vars, y, y_pred
                )
            logs = metric.stateless_result(this_metric_vars)
            new_metrics_vars += this_metric_vars

        
        state = (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            new_metrics_vars,
        )
        return logs, state





inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=3)





class CustomModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
        self.loss_fn = keras.losses.MeanSquaredError()

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

    def train_step(self, state, data):
        (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        ) = state
        x, y = data

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

        
        (loss, (y_pred, non_trainable_variables)), grads = grad_fn(
            trainable_variables,
            non_trainable_variables,
            x,
            y,
            training=True,
        )

        
        (
            trainable_variables,
            optimizer_variables,
        ) = self.optimizer.stateless_apply(
            optimizer_variables, grads, trainable_variables
        )

        
        loss_tracker_vars = metrics_variables[
            : len(self.loss_tracker.variables)
        ]
        mae_metric_vars = metrics_variables[len(self.loss_tracker.variables) :]

        loss_tracker_vars = self.loss_tracker.stateless_update_state(
            loss_tracker_vars, loss
        )
        mae_metric_vars = self.mae_metric.stateless_update_state(
            mae_metric_vars, y, y_pred
        )

        logs = {}
        logs[self.loss_tracker.name] = self.loss_tracker.stateless_result(
            loss_tracker_vars
        )
        logs[self.mae_metric.name] = self.mae_metric.stateless_result(
            mae_metric_vars
        )

        new_metrics_vars = loss_tracker_vars + mae_metric_vars

        
        state = (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            new_metrics_vars,
        )
        return logs, state

    @property
    def metrics(self):
        
        
        
        return [self.loss_tracker, self.mae_metric]



inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)


model.compile(optimizer="adam")


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)





class CustomModel(keras.Model):
    def test_step(self, state, data):
        
        x, y = data
        (
            trainable_variables,
            non_trainable_variables,
            metrics_variables,
        ) = state

        
        y_pred, non_trainable_variables = self.stateless_call(
            trainable_variables,
            non_trainable_variables,
            x,
            training=False,
        )
        loss = self.compute_loss(x, y, y_pred)

        
        new_metrics_vars = []
        for metric in self.metrics:
            this_metric_vars = metrics_variables[
                len(new_metrics_vars) : len(new_metrics_vars)
                + len(metric.variables)
            ]
            if metric.name == "loss":
                this_metric_vars = metric.stateless_update_state(
                    this_metric_vars, loss
                )
            else:
                this_metric_vars = metric.stateless_update_state(
                    this_metric_vars, y, y_pred
                )
            logs = metric.stateless_result(this_metric_vars)
            new_metrics_vars += this_metric_vars

        
        state = (
            trainable_variables,
            non_trainable_variables,
            new_metrics_vars,
        )
        return logs, state



inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(loss="mse", metrics=["mae"])


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.evaluate(x, y)