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

Author: Google for Developers - Bewertung: 7x - Views:72

Do you want the speed and functional power of JAX without losing the high-level convenience of model.fit? In this video, Google ML Developer Advocate Yufeng Guo (@yufengg) explains how Keras implements the principle of Progressive Disclosure of Complexity.



Learn how to take full control of your learning algorithms by overriding the train_step() and test_step() methods while keeping access to built-in callbacks, distribution support, and evaluation tools.



What You’ll Learn:

- Why override train_step instead of writing a loop from scratch?

- Understanding how to handle trainable variables, non-trainable variables, and optimizer states in a functional environment.

- Creating a compute_loss_and_updates function to manage forward passes and auxiliary data.

- Using jax.value_and_grad to compute gradients and losses simultaneously.

- Updating evaluation metrics using stateless_update_state.





Chapters:

0:00 - Introduction & The Default model.fit()

0:18 - Customizing Keras Training Loops

0:46 - Overriding train_step()

1:14 - Setting up the JAX Backend

1:26 - The Stateless train_step

2:11 - Stateless Loss Computation

3:04 - Taking Gradients in train_step

4:06 - How to pass around non-trainable variables

4:43 - Updating the Model Weights

5:05 - Handling Metrics

5:21 - Custom Evaluation Loops (overriding test_step)





Resources:

Complete Code Example →https://goo.gle/4eeSvlD

Keras Documentation → https://goo.gle/42Ebpv0

Keras Developer Guides →https://goo.gle/4um97N3



Subscribe to Google for Developers → https://goo.gle/developers



Speaker: Yufeng Guo

Products Mentioned: Google AI