TensorFlow with Docker

Code Examples and Tutorial

How Can You Use TensorFlow with Docker?

TensorFlow is an open source framework for running machine learning (ML), deep learning, and various statistical and predictive analytics workloads. Data scientists, predictive modelers, and statisticians use TensorFlow to streamline the development and execution of advanced analytics applications.

Docker is an open source platform for creating, deploying, and managing virtualized application containers on a shared operating system using an ecosystem of tools. You can use Docker containers to create a virtual environment that isolates a TensorFlow installation and replicate it to train and run models from any location.

TensorFlow Docker Requirements

The Docker platform enables you to build, deliver, and operate applications. You can use it to decouple applications from the infrastructure to quickly release software and manage the infrastructure like you manage applications. Docker's approach for shipping, testing, and deploying code can help minimize the time-to-release.

To start using Docker and TensorFlow, you need to install Docker on your local host machine and install NVIDIA Docker Support for GPU Support on Linux. You can leverage the NVIDIA Container Toolkit to create and execute GPU-accelerated Docker containers. The toolkit includes a container runtime library and utilities.

There is no need to install the CUDA Toolkit on the host system. However, you must install the NVIDIA driver on the host. Once the NVIDIA driver and Docker engine are installed on your Linux distribution, you can get started.

You can find additional installation instructions on the Docker website:

Getting Started with TensorFlow Docker Container

Deploying TensorFlow in a Docker Container

Once you’ve installed Docker, you can continue with the following steps:

  1. Open a terminal and run the following command to pull the TensorFlow Docker image:
    docker pull tensorflow/tensorflow:latest-gpu

    This will download the latest version of the TensorFlow Docker image with GPU support. If you don't have a GPU or don't need GPU support, you can use the following command to pull the CPU-only version of the image:
    docker pull tensorflow/tensorflow:latest
  2. Once the TensorFlow Docker image has been downloaded, you can use the docker run command to start a new container instance based on the image. For example, to start a new TensorFlow container with a Jupyter notebook server, you can use the following command:
    docker run -it --rm -p 8888:8888 tensorflow/tensorflow:latest-gpu

    This will start a new container in interactive mode (-i), remove it when it exits (-t), and expose the Jupyter notebook server on port 8888 (-p 8888:8888). If you want to use the CPU-only version of the image, you can use the following command instead:
    docker run -it --rm -p 8888:8888 tensorflow/tensorflow:latest
  3. Once the container is running, you can access the Jupyter notebook server by opening a web browser and going to the URL http://localhost:8888. You should see the Jupyter homepage, which will allow you to create and access notebooks.
  4. You can now start working with TensorFlow in your Jupyter notebooks. You can use the TensorFlow API to build and train machine learning models, and you can use the Jupyter notebook environment to document and share your work.

CPU-Only Image Example

To use TensorFlow with a CPU-only image, you can pull the TensorFlow Docker image from the Docker Hub using the following command:

docker pull tensorflow/tensorflow:latest

This will download the latest version of the TensorFlow Docker image that is optimized for use with CPUs. Once the image has been downloaded, you can start a new container instance using the docker run command:

docker run -it --rm -p 8888:8888 tensorflow/tensorflow:latest

This will start a new container in interactive mode (-i), remove it when it exits (-t), and expose the Jupyter notebook server on port 8888 (-p 8888:8888). You can then access the Jupyter notebook server by going to the URL http://localhost:8888 in a web browser.

To use TensorFlow with a GPU-enabled image, you will need to have a GPU available on your system and install the appropriate drivers and libraries. You can then pull the TensorFlow Docker image with GPU support using the following command:

docker pull tensorflow/tensorflow:latest-gpu

This will download the latest version of the TensorFlow Docker image that is optimized for use with GPUs. You can then start a new container instance using the docker run command:

docker run -it --rm -p 8888:8888 tensorflow/tensorflow:latest-gpu

This will start a new container in interactive mode (-i), remove it when it exits (-t), and expose the Jupyter notebook server on port 8888 (-p 8888:8888). You can then access the Jupyter notebook server by going to the URL http://localhost:8888 in a web browser.

GPU-Enabled Image Example

To use TensorFlow with GPU support, you will need to use a GPU-enabled image and ensure that your TensorFlow code is written to take advantage of the GPU. You can use the TensorFlow device function to specify which device (CPU or GPU) you want to use for a particular operation. For example:

import tensorflow as tf


with tf.device('/cpu:0'):

 # Code to run on the CPU


with tf.device('/gpu:0'):

 # Code to run on the GPU


You can also use the tf.config.experimental.set_visible_devices function to specify which GPUs are available to TensorFlow:

import tensorflow as tf


# Make GPU 0 visible to TensorFlow

tf.config.experimental.set_visible_devices(['/gpu:0'], 'GPU')

Quick Tutorial: TensorFlow Serving with Docker

TensorFlow Serving is a tool that allows you to serve machine learning models in a production environment. It is designed to be flexible, scalable, and easy to use, and it can be used with a variety of platforms and architectures, including Docker.

Step 1: Export Model Using model.save

  1. Before you can serve your model with TensorFlow Serving, you will need to export it in a format that TensorFlow Serving can understand. You can use the tf.saved_model.save function to save your model in a SavedModel format, which can be loaded by TensorFlow Serving.
    import tensorflow as tf

    # Build and compile your TensorFlow model
    load_data()
    pre_process_data()
    model = ...
    model.compile(...)

    # Save the model in a SavedModel format
    tf.saved_model.save(model, 'path/to/saved/model')

  2. Once you have exported your model, you can use the docker run command to start a new TensorFlow Serving container. You will need to mount the directory containing your SavedModel as a volume and expose the TensorFlow Serving port (8500) to allow clients to connect to the server. For example:
    docker run -it --rm -p 8500:8500 -v /path/to/saved/model:/models/my_model tensorflow/serving:latest

    This will start a new TensorFlow Serving container in interactive mode (-i), remove it when it exits (-t), and expose the TensorFlow Serving port on port 8500 (-p 8500:8500). It will also mount the directory containing your SavedModel as a volume (-v /path/to/saved/model:/models/my_model) so that TensorFlow Serving can access it.
  3. Test the TensorFlow Serving server: Once the TensorFlow Serving container is running, you can use a client to send requests to the server and test that it is working correctly. You can use the tensorflow_model_server command to start a client and send requests to the server. For example:
    tensorflow_model_server --model_name=my_model --model_base_path=/models/my_model

    This will start a client that connects to the TensorFlow Serving server and sends requests to the my_model model. You can then use the client to send requests to the server and test that it is working correctly.

Step 2: Pulling and Running a Serving Image

Now you can do the following:

  1. Open a terminal and use the docker pull command to download the serving image from the Docker Hub or a private registry. For example:
    docker pull my_image:latest

    This will download the latest version of the my_image serving image.
  2. Once the serving image has been downloaded, you can use the docker run command to start a new serving container based on the image. For example:
    docker run -it --rm -p 8080:8080 my_image:latest

    This will start a new serving container in interactive mode (-i), remove it when it exits (-t), and expose the serving port on port 8080 (-p 8080:8080). You can then access the serving service by connecting to the container's exposed port (e.g., http://localhost:8080).

    You can also use the docker run command to specify additional options, such as mounting volumes, setting environment variables, and linking to other containers. For example:
    docker run -it --rm -p 8080:8080 -v /path/to/data:/data -e MY_VAR=value --link db_container:db my_image:latest

    This will start a new serving container in interactive mode (-i), remove it when it exits (-t), and expose the serving port on port 8080 (-p 8080:8080). It will also mount the /path/to/data directory as a volume (-v /path/to/data:/data), set the MY_VAR environment variable (-e MY_VAR=value), and link to the db_container container as db (--link db_container:db).

Step 3: Serving with GPU

To use the GPU in your serving code, you will need to ensure that your code is written to take advantage of the GPU. This may include using GPU-accelerated libraries and frameworks, such as TensorFlow or PyTorch, and specifying which device (CPU or GPU) you want to use for particular operations.

For example, in TensorFlow you can use the device function to specify which device you want to use for a particular operation:

import tensorflow as tf


with tf.device('/cpu:0'):

 # Code to run on the CPU


with tf.device('/gpu:0'):

 # Code to run on the GPU


You can also use the tf.config.experimental.set_visible_devices function to specify which GPUs are available to TensorFlow:

import tensorflow as tf


# Make GPU 0 visible to T

TensorFlow Virtualization with Run:ai

Run:ai automates resource management and workload orchestration for machine learning infrastructure. With Run:ai, you can automatically run as many compute intensive experiments as needed in TensorFlow and other deep learning frameworks.

Here are some of the capabilities you gain when using Run:ai:

  • Advanced visibility: create an efficient pipeline of resource sharing by pooling GPU compute resources.
  • No more bottlenecks: you can set up guaranteed quotas of GPU resources, to avoid bottlenecks and optimize billing.
  • A higher level of control: Run:ai enables you to dynamically change resource allocation, ensuring each job gets the resources it needs at any given time.

Run:ai simplifies machine learning infrastructure pipelines, helping data scientists accelerate their productivity and the quality of their models.

Learn more about the Run:ai platform