Search⌘ K
AI Features
Home
Courses
Deep Learning with JAX and Flax

Solution: Optimizers

Explore how to load and preprocess the Fashion MNIST dataset, define a vanilla CNN model using Flax's linen module, and implement training with AdaBelief and SGD optimizers. Understand batch training, loss calculation, and accuracy evaluation to improve model performance using JAX and Flax.

We'll cover the following...

Loading dataset

We use the keras.dataset library to load and visualize the Fashion MNIST dataset.

Python 3.8
from tensorflow import keras
from sklearn.model_selection import train_test_split
from jax import numpy as jnp
import numpy as np
from matplotlib import pyplot
(X_train, Y_train), (X_test, Y_test) = keras.datasets.fashion_mnist.load_data()
X = np.concatenate((X_train, X_test))
Y = np.concatenate((Y_train, Y_test))
train_size = 0.8
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, train_size=0.04, test_size=0.01, random_state=2019)
for i in range(9):
    pyplot.subplot(330 + 1 + i)
    pyplot.imshow(X_train[i])
pyplot.savefig("output/image.png")
X_train, X_test, Y_train, Y_test = jnp.array(X_train, dtype=jnp.float32),\
                                   jnp.array(X_test, dtype=jnp.float32),\
                                   jnp.array(Y_train, dtype=jnp.float32),\
                                   jnp.array(Y_test, dtype=jnp.float32)
X_train, X_test = X_train.reshape(-1,28,28,1), X_test.reshape(-1,28,28,1)
X_train, X_test = X_train/255.0, X_test/255.0
classes =  jnp.unique(Y_train)
print('X_train', X_train.shape, 'X_test',X_test.shape,)   
print('Y_train',Y_train.shape, 'Y_test',Y_test.shape)

In the code above:

  • Lines 1–5: We import the keras library from TensorFlow to load the dataset, the train_test_split method from sklearn.mode_selection to split the dataset, and the JAX version of NumPy to perform numerical operations. Also, we import the numpy and matplotlib libraries for visualization.

  • Lines 7–9: We load the Fashion MNIST dataset and combine the training and test datasets.

  • Lines 11–12: We define the train_size as 0.8 and split the train and test dataset.

  • Lines 14–16: We use the for loop to create a plot of 9 images to display in the output.

  • Line 17: We save the image in the output folder to show in the output of the playground.

  • Lines 19–23: We convert the train and test data to the JAX arrays.

  • Line 24: We reshape the dataset with the ...