Deep Learning with JAX and Flax/

...

/

Solution: Optimizers

Solution: Optimizers

Let’s review the solution.

We'll cover the following...

Loading dataset

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

Press + to interact
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 ...