ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

TensorFlow 2 quickstart for beginners

View on TensorFlow.org View source on GitHub Download notebook

This short introduction uses Keras to:

  1. Load a prebuilt dataset.
  2. Build a neural network machine learning model that classifies images.
  3. Train this neural network.
  4. Evaluate the accuracy of the model.

This tutorial is a Google Colaboratory notebook. Python programs are run directly in the browser—a great way to learn and use TensorFlow. To follow this tutorial, run the notebook in Google Colab by clicking the button at the top of this page.

  1. In Colab, connect to a Python runtime: At the top-right of the menu bar, select CONNECT.
  2. Run all the notebook code cells: Select Runtime > Run all.

Set up TensorFlow

Import TensorFlow into your program to get started:

import tensorflow as tf
print("TensorFlow version:", tf.__version__)
TensorFlow version: 2.6.0

If you are following along in your own development environment, rather than Colab, see the install guide for setting up TensorFlow for development.

Load a dataset

Load and prepare the MNIST dataset. Convert the sample data from integers to floating-point numbers:

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step

Build a machine learning model

Build a tf.keras.Sequential model by stacking layers.

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

For each example, the model returns a vector of logits or log-odds scores, one for each class.

predictions = model(x_train[:1]).numpy()
predictions
array([[ 0.74148655, -0.39261633,  0.08016336, -0.46431944,  0.21458861,
         0.31183302,  0.7555975 ,  0.80728006, -0.6296631 , -0.4926056 ]],
      dtype=float32)

The tf.nn.softmax function converts these logits to probabilities for each class:

tf.nn.softmax(predictions).numpy()
array([[0.16651046, 0.05356818, 0.08594736, 0.04986165, 0.09831339,
        0.10835411, 0.16887674, 0.1778342 , 0.04226285, 0.04847102]],
      dtype=float32)

Define a loss function for training using losses.SparseCategoricalCrossentropy, which takes a vector of logits and a True index and returns a scalar loss for each example.

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

This loss is equal to the negative log probability of the true class: The loss is zero if the model is sure of the correct class.

This untrained model gives probabilities close to random (1/10 for each class), so the initial loss should be close to -tf.math.log(1/10) ~= 2.3.

loss_fn(y_train[:1], predictions).numpy()
2.2223506

Before you start training, configure and compile the model using Keras Model.compile. Set the optimizer class to adam, set the loss to the loss_fn function you defined earlier, and specify a metric to be evaluated for the model by setting the metrics parameter to accuracy.

model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

Train and evaluate your model

Use the Model.fit method to adjust your model parameters and minimize the loss:

model.fit(x_train, y_train, epochs=5)
Epoch 1/5
1875/1875 [==============================] - 3s 2ms/step - loss: 0.2956 - accuracy: 0.9137
Epoch 2/5
1875/1875 [==============================] - 3s 2ms/step - loss: 0.1430 - accuracy: 0.9572
Epoch 3/5
1875/1875 [==============================] - 3s 2ms/step - loss: 0.1059 - accuracy: 0.9682
Epoch 4/5
1875/1875 [==============================] - 3s 2ms/step - loss: 0.0870 - accuracy: 0.9736
Epoch 5/5
1875/1875 [==============================] - 3s 2ms/step - loss: 0.0732 - accuracy: 0.9774
<keras.callbacks.History at 0x7fabfe6f6cd0>

The Model.evaluate method checks the models performance, usually on a "Validation-set" or "Test-set".

model.evaluate(x_test,  y_test, verbose=2)
313/313 - 0s - loss: 0.0771 - accuracy: 0.9756
[0.07707860320806503, 0.975600004196167]

The image classifier is now trained to ~98% accuracy on this dataset. To learn more, read the TensorFlow tutorials.

If you want your model to return a probability, you can wrap the trained model, and attach the softmax to it:

probability_model = tf.keras.Sequential([
  model,
  tf.keras.layers.Softmax()
])
probability_model(x_test[:5])
<tf.Tensor: shape=(5, 10), dtype=float32, numpy=
array([[9.35628634e-07, 1.90160367e-08, 2.08603578e-05, 3.36712168e-04,
        3.99928729e-10, 3.29815066e-07, 1.11730861e-12, 9.99547541e-01,
        5.78590800e-07, 9.29015441e-05],
       [9.06768847e-08, 1.75887300e-03, 9.98239040e-01, 1.42416525e-06,
        2.24498985e-14, 1.93019602e-07, 4.19461571e-07, 1.78353019e-14,
        1.09490275e-08, 3.30677152e-15],
       [1.80358853e-07, 9.99705374e-01, 4.36524460e-05, 7.87066620e-06,
        6.60410151e-05, 1.29832119e-06, 2.66996949e-06, 5.51502962e-05,
        1.15344104e-04, 2.45894512e-06],
       [9.99937892e-01, 1.82150364e-10, 2.67099913e-05, 1.87603248e-07,
        2.81131548e-07, 1.86480634e-06, 2.33697901e-05, 1.97653662e-07,
        3.29448504e-08, 9.38197445e-06],
       [1.53162068e-06, 2.57238808e-09, 2.05853598e-06, 2.75526531e-07,
        9.93477166e-01, 2.57349512e-07, 3.59538944e-06, 9.20545644e-05,
        1.77590653e-06, 6.42123120e-03]], dtype=float32)>

Conclusion

Congratulations! You have trained a machine learning model using a prebuilt dataset using the Keras API.

For more examples of using Keras, check out the tutorials. To learn more about building models with Keras, read the guides. If you want learn more about loading and preparing data, see the tutorials on image data loading or CSV data loading.