Using TensorFlow in Python A Comprehensive Guide

TensorFlow is an open-source deep learning library developed by the Google Brain team. It is designed to facilitate the creation, training, and deployment of machine learning models with a focus on neural networks. In this article, we’ll explore how to use TensorFlow in Python, covering essential topics like data preparation, model creation, training, evaluation, and deployment.
Table of Contents
Section titled “Table of Contents”- Prerequisites
- Installation
- Data Preparation
- Building a Neural Network Model
- Training the Model
- Evaluating the Model
- Saving and Loading Models
- Deploying the Model
- Conclusion
1. Prerequisites
Section titled “1. Prerequisites”To follow this guide, you should have a basic understanding of Python programming and familiarity with machine learning concepts. Knowledge of neural networks and deep learning is helpful but not required.
2. Installation
Section titled “2. Installation”To install TensorFlow, you can use the pip package manager. It is recommended to use a virtual environment to avoid conflicts with other packages. Run the following command:
pip install tensorflowFor GPU support, install the GPU version of TensorFlow:
pip install tensorflow-gpuEnsure that you have the appropriate CUDA and cuDNN libraries installed on your system for GPU support.
3. Data Preparation
Section titled “3. Data Preparation”Before building a model, we need to prepare the data. Let’s consider the popular MNIST dataset for handwritten digit recognition. We can load this dataset using the TensorFlow Datasets module:
import tensorflow as tfimport tensorflow_datasets as tfds
(ds_train, ds_test), ds_info = tfds.load( 'mnist', split=['train', 'test'], shuffle_files=True, as_supervised=True, with_info=True,)
## Preprocessingdef normalize_img(image, label): return tf.cast(image, tf.float32) / 255., label
ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)ds_test = ds_test.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
## Batch and prefetchds_train = ds_train.batch(128).prefetch(tf.data.experimental.AUTOTUNE)ds_test = ds_test.batch(128).prefetch(tf.data.experimental.AUTOTUNE)In this code snippet, we load the MNIST dataset, normalize the images by dividing pixel values by 255, and create batches of size 128 for training and testing. The prefetch method is used to optimize data loading performance.
4. Building a Neural Network Model
Section titled “4. Building a Neural Network Model”Now that we’ve prepared our data, let’s build a simple neural network model using TensorFlow’s Keras API:
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, activation='softmax')])
model.compile( optimizer=tf.keras.optimizers.Adam(0.001), loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],)This code snippet defines a simple feedforward neural network with one hidden layer containing 128 neurons and a dropout layer for regularization. The output layer has 10 neurons with softmax activation, representing the probability distribution over the 10 digit classes.
5. Training the Model
Section titled “5. Training the Model”With our model defined, we can now train it using the fit method:
history = model.fit( ds_train, epochs=10, validation_data=ds_test,)This code snippet trains the model for 10 epochs using the training data and validates it using the test data. The training history is stored in the history variable, which can be used to plot training progress.
6. Evaluating the Model
Section titled “6. Evaluating the Model”To evaluate the model on the test dataset, we can use the evaluate method:
test_loss, test_accuracy = model.evaluate(ds_test)print('Test loss:', test_loss)print('Test accuracy:', test_accuracy)This code snippet calculates the test loss and accuracy, providing a measure of the model’s performance on unseen data.
7. Saving and LoadingModels
Section titled “7. Saving and LoadingModels”After training and evaluating our model, we might want to save it for future use. TensorFlow allows us to save models in the HDF5 format or the TensorFlow SavedModel format. Here’s how to save and load a model in both formats:
Saving and Loading in HDF5 Format
Section titled “Saving and Loading in HDF5 Format”## Save the modelmodel.save('mnist_model.h5')
## Load the saved modelloaded_model = tf.keras.models.load_model('mnist_model.h5')Saving and Loading in TensorFlow SavedModel Format
Section titled “Saving and Loading in TensorFlow SavedModel Format”## Save the modelmodel.save('mnist_saved_model')
## Load the saved modelloaded_model = tf.keras.models.load_model('mnist_saved_model')By saving the model, we can later load it to make predictions, continue training, or deploy it to production.
8. Deploying the Model
Section titled “8. Deploying the Model”Once our model is trained, evaluated, and saved, we may want to deploy it to a production environment. TensorFlow Serving is a flexible, high-performance serving system for machine learning models designed for production environments. It supports both TensorFlow and non-TensorFlow machine learning models.
To deploy a model using TensorFlow Serving, we first need to install it on our system. Follow the official installation instructions to set up TensorFlow Serving.
After installing TensorFlow Serving, we can deploy our model:
- Start the TensorFlow Serving server by pointing it to the SavedModel directory:
tensorflow_model_server --port=8501 --rest_api_port=8502 --model_name=mnist --model_base_path=/path/to/mnist_saved_model- Now, we can make a prediction using our deployed model by sending a REST API request:
import requestsimport jsonimport numpy as np
## Load a sample imagesample_image, sample_label = ds_test.take(1).as_numpy_iterator().next()sample_image = sample_image.reshape(1, 28, 28)
## Prepare the request payloaddata = json.dumps({"signature_name": "serving_default", "instances": sample_image.tolist()})
## Send the requestheaders = {"content-type": "application/json"}response = requests.post('http://localhost:8502/v1/models/mnist:predict', data=data, headers=headers)
## Parse the responsepredictions = json.loads(response.text)['predictions']predicted_label = np.argmax(predictions[0])
print('Predicted label:', predicted_label)print('Actual label:', sample_label)In this code snippet, we load a sample image from the test dataset, prepare a JSON payload, and send a prediction request to the deployed TensorFlow Serving server. Then, we parse the response and print the predicted and actual labels.
9. Conclusion
Section titled “9. Conclusion”In this article, we covered the core concepts and steps to use TensorFlow in Python, including data preparation, model creation, training, evaluation, and deployment. TensorFlow is a powerful and versatile library that enables developers to build, train, and deploy machine learning models with ease. By mastering TensorFlow, you can unlock the potential of deep learning and create groundbreaking applications in fields like computer vision, natural language processing, and reinforcement learning.
© 2026 David Li · Home · Blog · RSS ·Classic site · Built with Astro, Starlight and three.js