Unlocking Low-Level Control: Customizing Keras Training Loops with JAX

0:00 When working with deep learning models,

0:02 Keras provides a convenient model.fit API for training.

0:07 But what if your algorithm requires a custom training procedure and you still

0:12 want to benefit from high level features

0:15 like callbacks and built in distribution support.

0:18 Well, Keras follows an important core

0:21 principle the progressive disclosure of complexity.

0:26 This means you should be able to gain more low

0:28 level control without having to abandon all that high level convenience.

0:35 Hi there, my name is Yufeng and today we will look at how to customize the Keras

0:40 training loop with the JAX backend by overriding

0:44 the training step to customize what fit does.

0:49 One approach is to override the trainstep function of the model class.

0:55 This is the function that is called by fit for every batch of data,

1:00 allowing you to run your own learning algorithm while still using fit.

1:05 As usual.

1:07 This approach works for sequential models.

1:10 Functional API models or subclass models.

1:14 Since this customization specifically targets JAX.

1:17 Optimization, we must ensure our environment is

1:21 correctly configured to use the JAX backend.

1:23 Before importing Keras.

1:26 Customizing the training step in JAX

1:28 requires understanding JAX's reliance on stateless computation.

1:34 The entire trainstep method must be fully stateless.

1:39 What does this mean.

1:40 Well, it means that all elements of the model state the trainable variables,

1:45 the non-trainable variables, optimizer variables,

1:48 and the metrics variables are all explicitly

1:52 passed as inputs within a state tuple,

1:56 and then returned as updated versions of those same variables.

2:01 We'll use stateless versions of the call apply and compute loss functions,

2:06 so keep an eye out for how their inputs and outputs end up being structured.

2:11 A key part of the JAKs implementation is calculating the gradients,

2:15 typically using a helper function.

2:17 In this example, we made one called compute loss and updates.

2:21 As the name suggests, it computes the loss and the updates

2:26 portion refers to updating those non-trainable variables.

2:30 So two things are happening in this function.

2:33 We first do a forward pass of the model using model call,

2:37 passing the trainable variables

2:39 and non-trainable variables explicitly as inputs.

2:42 Using stateless call is important with the JAX backend,

2:47 and this method will return the predictions y pred,

2:52 and any updated non-trainable variables as well.

2:55 Then we'll use these values to compute the loss by passing

2:58 in the expected y values along with those predicted y values.

3:04 Now let's look at the train step where

3:06 we call our helper function compute loss and updates.

3:10 The plan is to use JAX's grad function transformation,

3:15 which will produce a new function

3:17 that computes the gradient of our helper function.

3:21 This means we'll end up with a function that computes the gradient of the loss,

3:26 which is perfect for our machine learning needs.

3:29 In practice, we want to compute the value of the loss as well as its gradient,

3:35 so we'll do that to our helper function using JAX value and grad.

3:42 This is simply combining the operations of calling

3:45 our helper function itself and creating its gradient function.

3:50 When we compare doing this manually

3:52 with separate calls to the function and grad, we'll see they're very similar,

3:58 except you'll notice we didn't have to write our list

4:01 of arguments to compute loss and updates twice, which is nice.

4:07 So you might be wondering also why we use the has x equals true argument here.

4:14 This argument indicates that the function being differentiated returns a pair

4:18 where the first element is the value to be differentiated,

4:22 and the second is auxiliary data that should not be differentiated.

4:28 Our loss computation function returns exactly this the loss,

4:32 which we want to differentiate,

4:34 followed by some other stuff that we're kind of just passing

4:37 through the widespread and those updated

4:40 non-trainable variables inside of trainstep.

4:44 After computing the gradients, we update the variables.

4:48 Keras provides stateless methods for this to.

4:51 To apply the gradients, we use the optimizers stateless apply method.

4:56 It's aptly named optimizer.zero grad apply,

5:00 which returns the updated trainable variables

5:02 along with the updated optimizer variables.

5:05 If you are using Keras metrics,

5:07 you can handle them by calculating the metric result

5:10 inside the step function using metrics stateless update state.

5:15 Ensuring these new metric variables are also returned in the state tuple.

5:21 This same capability for low level control can be applied to evaluation.

5:27 So to customize how model.evaluate operates,

5:31 override the test set method with a call to compute loss and updates,

5:35 just like we did in trainstep.

5:37 Except we don't need to perform the stateless

5:40 apply to update the weights this time,

5:42 since we're just evaluating and we can move

5:45 directly to recording our loss metrics for evaluation.

5:48 Scoring by subclassing Keras model and overriding

5:54 methods like trainstep and test step,

5:56 you gain fine grained control over the training

5:59 and evaluation algorithms using JAX's powerful function tools,

6:04 all while preserving the convenient high level features provided by model.fit.

6:10 So what custom algorithms have you integrated into Keras model.

6:14 Share your thoughts in the comments below.

6:16 And for the full documentation and complete code examples,

6:19 check out the guide LinkedIn the description below.

6:22 Thanks for watching and I'll catch you in the next one.