Gain insights into JAX and Flax's features for deep learning. Learn about optimizers, functions, data loading, and model training. Explore hands-on projects for practical experience.

edujax_1.tar.gz

jaxbook

jax-code-widget

This course comprehensively introduces JAX and Flax, two open-source libraries that have gained prominence for their efficiency, flexibility, and scalability in deep learning applications. 

In this course, you’ll explore deep learning principles and understand the unique features of JAX and Flax. You will learn the basics of JAX, optimizers using JAX and Flax, and loss and activation functions. You’ll also learn how to load datasets, perform classification using distributed learning, and use ResNet and LSTM models. In the end, you will complete a project for hands-on experience using JAX and Flax for transfer learning. 

By the end of the course, you’ll be proficient in implementing and customizing neural network models using JAX and Flax, equipped with hands-on skills in advanced optimization and distributed training.

Deep Learning with JAX and Flax

Tracking machine learning experiments makes understanding and visualizing the model's performance easy. It also makes it possible to spot any problems in the network. For example, we can quickly spot overfitting by looking at the training and validation charts. We can plot these charts using our favorite charts package, such as Matplotlib. However, we can also use more advanced tools, such as TensorBoard.

**TensorBoard** is an open-source library that provides tools for experiment tracking in machine learning. We can use TensorBoard for:

- Tracking and visualizing model evaluation metrics, such as accuracy.
- Logging images.
- Visualizing hyper-parameter tuning.
- Project embeddings such as word embedding in natural language processing problems.
- Visualizing histograms of the model's weights and biases.
- Plotting the architecture of the model.
- Profiling the performance of the network.

We can use TensorBoard with popular machine learning libraries, such as XGBoost, JAX, Flax, and PyTorch.

This lesson will focus on how to use TensorBoard when building networks with JAX and Flax.

# Using TensorBoard with Jupyter Notebooks and Google Colaboratory

We can load in our environment, usually Google Colaboratory or the local notebook.

```python
%load_ext tensorboard
```

Next, inform TensorBoard which folder will contain the log information.

```python
log_folder = "runs"
```

# Launching TensorBoard

Tensorboard is launched using the `tensorboard` magic command in notebook environments while specifying the `logdir`.

```python
%tensorboard --logdir={log_folder}
```

# TensorBoard dashboards

TensorBoard has various dashboards for showing different types of information:

- The **Scalars dashboard** tracks numerical information, such as training metrics per epoch. We can use it to track other scalar values, such as model training speed and learning rate.





Learn how to use TensorBoard with JAX and Flax to train models.

Course Introduction

Basics of JAX

Optimizers in JAX and Flax

Loss and Activation Functions

Load Datasets in JAX

Image Classification and Distributed Training

TensorBoard and State Handling

LSTM in JAX and Flax

Flax vs. TensorFlow

Using ResNet Model in Flax

Transfer Learning in JAX and Flax

Conclusion

Appendix

Setting Up TensorBoard