JAX - Machine learning nuggets

Sign in Subscribe

JAX

A collection of 15 posts

Convolutional Neural Networks in JAX: Ultimate Guide

Convolutional Neural Networks in JAX: Ultimate Guide

JAX is a high performance library that offers accelerated computing through XLA and Just In Time Compilation. It also has handy features that enable you to write one codebase that can be applied to batches of data and run on CPU, GPU, or TPU. However, one of its biggest selling

Train ResNet in Flax from scratch(Distributed ResNet training)

Train ResNet in Flax from scratch(Distributed ResNet training)

Apart from designing custom CNN architectures, you can use architectures that have already been built. ResNet is one such popular architecture. In most cases, you'll achieve better performance by using such architectures. In this article, you will learn how to perform distributed training of a ResNet model in

Handling state in JAX & Flax (BatchNorm and DropOut layers)

Handling state in JAX & Flax (BatchNorm and DropOut layers)

Jitting functions in Flax makes them faster but requires that the functions have no side effects. The fact that jitted functions can't have side effects introduces a challenge when dealing with stateful items such as model parameters and stateful layers such as batch normalization layers. In this article,

Transfer learning with JAX & Flax

Transfer learning with JAX & Flax

Training large neural networks can take days or weeks. Once these networks are trained, you can take advantage of their weights and apply them to new tasks– transfer learning. As a result, you fine-tune a new network and get good results in a short period. Let's look at

Flax vs. TensorFlow

Flax vs. TensorFlow

Flax is the neural network library for JAX. TensorFlow is a deep learning library with a large ecosystem of tools and resources. Flax and TensorFlow are similar but different in some ways. For instance, both Flax and TensorFlow can run on XLA. Let's look at the differences between

Activation functions in JAX and Flax

Activation functions in JAX and Flax

Activation functions are applied in neural networks to ensure that the network outputs the desired result. The activations functions cap the output within a specific range. For instance, when solving a binary classification problem, the outcome should be a number between 0 and 1. This indicates the probability of an

Optimizers in JAX and Flax

Optimizers in JAX and Flax

Optimizers are applied when training neural networks to reduce the error between the true and predicted values. This optimization is done via gradient descent. Gradient descent adjusts errors in the network through a cost function. In JAX, optimizers are applied from the Optax library. Optimizers can be classified into two

Elegy(High-level API for deep learning in JAX & Flax)

Elegy(High-level API for deep learning in JAX & Flax)

Training deep learning networks in Flax is done in a couple of steps. It involves creating the following functions: * Model definition. * Compute metrics. * Training state. * Training step. * Training and evaluation function. Flax and JAX give more control in defining and training deep learning networks. However, this comes with more verbosity.

LSTM in JAX & Flax (Complete example with code and notebook)

LSTM in JAX & Flax (Complete example with code and notebook)

LSTMs are a class of neural networks used to solve sequence problems such as time series and natural language processing. The LSTMs maintain some internal state that is useful in solving these problems. LSTMs apply for loops to iterate over each time step. We can use functions from JAX and

How to use TensorBoard in Flax

How to use TensorBoard in JAX & 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, you can quickly spot overfitting by looking at the training and validation charts. You can plot these charts using your favorite charts

How to load datasets in JAX with TensorFlow

How to load datasets in JAX with TensorFlow

JAX doesn't ship with data loading utilities. This keeps JAX focused on providing a fast tool for building and training machine learning models. Loading data in JAX is done using either TensorFlow or PyTorch. In the Image classification with JAX & Flax tutorial, we saw how to load

JAX loss functions

JAX loss functions

Loss functions are at the core of training machine learning. They can be used to identify how well the model is performing on a dataset. Poor performance leads to a very high loss, while a well-performing model will have a lower loss. Therefore, the choice of a loss function is

JAX (What it is and how to use it in Python)

JAX is a Python library offering high performance in machine learning with XLA and Just In Time (JIT) compilation. Its API is similar to NumPy's with a few differences. JAX ships with functionalities that aim to improve and increase speed in machine learning research. These functionalities include: * Automatic

Distributed training with JAX & Flax

Distributed training with JAX & Flax

Training models on accelerators with JAX and Flax differs slightly from training with CPU. For instance, the data needs to be replicated in the different devices when using multiple accelerators. After that, we need to execute the training on multiple devices and aggregate the results. Flax supports TPU and GPU

Image classification with JAX & Flax

Image classification with JAX & Flax

Flax is a neural network library for JAX. JAX is a Python library that provides high-performance computing in machine learning research. JAX provides an API similar to NumPy making it easy to adopt. JAX also includes other functionalities for improving machine learning research. They include: * Automatic differentiation. JAX supports forward