tf_nightly_keras_mnist_digit_classification_quickstart.py

1# Install tf-nightly if not already installed
2# !pip install tf-nightly
3
4import tensorflow as tf
5print("TensorFlow version:", tf.__version__)
6
7# Load and prepare the MNIST dataset
8mnist = tf.keras.datasets.mnist
9
10(x_train, y_train), (x_test, y_test) = mnist.load_data()
11x_train, x_test = x_train / 255.0, x_test / 255.0
12
13# Build the tf.keras.Sequential model by stacking layers
14model = tf.keras.models.Sequential([
15  tf.keras.layers.Flatten(input_shape=(28, 28)),
16  tf.keras.layers.Dense(128, activation='relu'),
17  tf.keras.layers.Dropout(0.2),
18  tf.keras.layers.Dense(10)
19])
20
21# Define the loss function
22loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
23
24# Compile the model
25model.compile(optimizer='adam',
26              loss=loss_fn,
27              metrics=['accuracy'])
28
29# Train the model
30model.fit(x_train, y_train, epochs=5)
31
32# Evaluate the model
33model.evaluate(x_test,  y_test, verbose=2)
34
35# If you want the model to return a probability, wrap the trained model and attach the softmax to it
36probability_model = tf.keras.Sequential([
37  model,
38  tf.keras.layers.Softmax()
39])
40
41print(probability_model(x_test[:5]))