Site Logo

TensorFlow with Docker: A Step-by-Step Guide to Containerized Machine Learning

Introduction

TensorFlow is a powerful open-source framework for machine learning, and running it inside Docker containers simplifies setup, ensures consistency, and enhances portability across environments. Docker allows you to package TensorFlow with its dependencies, making it easy to develop, test, and deploy models for projects like MNIST Classification, Face Recognition, or Scalable API. Whether you're a beginner or an experienced developer, using TensorFlow with Docker streamlines your workflow.

This guide provides a clear, replicable introduction to running TensorFlow in Docker, assuming no prior knowledge of containers. We’ll set up a Docker environment, run a TensorFlow container, and train a simple convolutional neural network (CNN) on the MNIST dataset. Each step explains what to do, why it matters, and how to do it, with a practical Dockerfile and script you can use to build and run your own container. By the end, you’ll be equipped to use TensorFlow with Docker for your own projects, like Stock Price Prediction or Real-Time Detection. This complements resources like What is TensorFlow?, TensorFlow Workflow, and TensorFlow in Deep Learning.

Step-by-Step Guide to Using TensorFlow with Docker

We’ll set up Docker, create a TensorFlow container, and train a CNN on the MNIST dataset (60,000 training and 10,000 test images of handwritten digits, 0–9, 28x28 pixels). This guide is designed for beginners, using a local environment (Windows, macOS, or Linux) and Docker Desktop. Each step is clear and replicable, ensuring you can run TensorFlow in a containerized environment.

Step 1: Install Docker

  • What You’re Doing: Setting up Docker on your computer.
  • Why It Matters: Docker is required to create and run containers, which package TensorFlow and its dependencies for consistency ([Installing TensorFlow](/tensorflow/introduction/installing-tensorflow)).
  • How to Do It:
  1. Download Docker Desktop from docker.com/get-started for your operating system (Windows, macOS, or Linux).
  2. Install Docker Desktop, following the prompts. On Windows, enable WSL 2 if prompted; on Linux, ensure Docker Engine is installed.
  3. Open a terminal (Command Prompt, PowerShell, or Terminal) and verify installation:
docker --version
Expect output like <mark>Docker version 20.x.x</mark>.
  1. Start Docker Desktop and ensure it’s running (check the system tray or docker info).
  • Tip: If you encounter issues, check Docker’s documentation ([docs.docker.com](https://docs.docker.com)) or restart Docker Desktop.

Step 2: Pull a TensorFlow Docker Image

  • What You’re Doing: Downloading an official TensorFlow Docker image from Docker Hub.
  • Why It Matters: TensorFlow images come pre-configured with TensorFlow and dependencies, saving setup time.
  • How to Do It:
  1. In your terminal, pull the latest TensorFlow image with GPU support (use cpu for non-GPU systems):
docker pull tensorflow/tensorflow:latest-gpu
For CPU-only:
docker pull tensorflow/tensorflow:latest
  1. Verify the image:
docker images
Look for <mark>tensorflow/tensorflow</mark> with the <mark>latest-gpu</mark> or <mark>latest</mark> tag.
  • Tip: GPU images require NVIDIA drivers and NVIDIA Container Toolkit; use CPU images for simplicity if you lack GPU setup ([TensorFlow GPU Support](/tensorflow/fundamentals/gpu-memory-optimization)).

Step 3: Run a TensorFlow Container

  • What You’re Doing: Starting a Docker container to run TensorFlow.
  • Why It Matters: The container provides an isolated environment for TensorFlow, ensuring consistency across machines.
  • How to Do It:
  1. Create a local directory for your project (e.g., tensorflow_project):
mkdir tensorflow_project
     cd tensorflow_project
  1. Run a TensorFlow container, mounting the directory for file access:
docker run -it -v $(pwd):/app -w /app tensorflow/tensorflow:latest bash
<ul>
<li><mark>-it</mark>: Interactive terminal.</li>
<li><mark>-v $(pwd):/app</mark>: Mounts your directory to <mark>/app</mark> in the container.</li>
<li><mark>-w /app</mark>: Sets the working directory to <mark>/app</mark>.</li>
<li><mark>bash</mark>: Opens a shell in the container.</li>
</ul>

3. Inside the container, verify TensorFlow:

python -c "import tensorflow as tf; print(tf.__version__)"
Expect output like <mark>2.16.2</mark>.
  • Tip: Use -p 8888:8888 to enable Jupyter notebooks (e.g., docker run -it -p 8888:8888 ...).

Step 4: Create a Custom Dockerfile

  • What You’re Doing: Building a custom Docker image with your TensorFlow script and dependencies.
  • Why It Matters: A custom image ensures reproducibility and portability for your project ([Docker for ML](/tensorflow/production/docker-for-ml)).
  • How to Do It:
  1. In your tensorflow_project directory, create a file named Dockerfile (no extension) with:
FROM tensorflow/tensorflow:latest
     WORKDIR /app
     COPY . /app
     RUN pip install matplotlib
     CMD ["python", "mnist.py"]
<ul>
<li><mark>FROM</mark>: Uses the TensorFlow CPU image.</li>
<li><mark>WORKDIR</mark>: Sets the working directory.</li>
<li><mark>COPY</mark>: Copies your files into the container.</li>
<li><mark>RUN</mark>: Installs additional dependencies (e.g., Matplotlib for visualization).</li>
<li><mark>CMD</mark>: Runs your script (<mark>mnist.py</mark>).</li>
</ul>

2. Create a Python script (mnist.py) for MNIST classification (see program below).

  • Tip: Use latest-gpu in the Dockerfile for GPU support if configured.

Step 5: Build and Run Your Custom Image

  • What You’re Doing: Building the Docker image and running it to train the model.
  • Why It Matters: This packages your code and TensorFlow, ensuring it runs consistently anywhere.
  • How to Do It:
  1. Build the image:
docker build -t my-tensorflow-app .
<ul>
<li><mark>-t my-tensorflow-app</mark>: Names the image.</li>
<li><mark>.</mark>: Uses the Dockerfile in the current directory.</li>
</ul>

2. Run the container:

docker run -v $(pwd):/app my-tensorflow-app
This executes <mark>mnist.py</mark>, training the model.
  1. Check the output for accuracy (~98–99%) and saved model.
  • Tip: Use -v to persist output files (e.g., model, logs) in your local directory.

Step 6: Monitor and Extend

  • What You’re Doing: Reviewing training results and exploring further Docker use.
  • Why It Matters: Monitoring ensures your model performs well, and Docker supports scaling to production ([Model Monitoring](/tensorflow/production/model-monitoring)).
  • How to Do It:
  1. Check TensorBoard logs (generated by the script) in your local directory:
tensorboard --logdir ./logs
Open the provided URL (e.g., <mark>http://localhost:6006</mark>) to view metrics.
  1. Extend your Dockerfile for other projects by adding dependencies or scripts.
  2. Explore Docker for production with TensorFlow Serving or cloud deployment (Cloud Integration).
  • Tip: Save models and logs to your local directory for persistence.

Practical Dockerfile and Script

Below is a Dockerfile and Python script to build and run a TensorFlow container for MNIST classification, demonstrating the steps above. For more examples, see TensorFlow in Deep Learning.

Prerequisites

  • Docker Desktop installed and running.
  • Local directory (tensorflow_project) with write permissions.
  • Terminal access (Command Prompt, PowerShell, or Terminal).

Dockerfile

Create Dockerfile in tensorflow_project:

FROM tensorflow/tensorflow:latest
WORKDIR /app
COPY . /app
RUN pip install matplotlib
CMD ["python", "mnist.py"]

Python Script

Create mnist.py in tensorflow_project:

import tensorflow as tf
import matplotlib.pyplot as plt

# Load and prepare MNIST data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# Create tf.data pipeline
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(60000).batch(32).prefetch(tf.data.AUTOTUNE)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32).prefetch(tf.data.AUTOTUNE)

# Build CNN model
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# Compile and train
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_ds, epochs=5, validation_data=test_ds,
          callbacks=[tf.keras.callbacks.TensorBoard(log_dir='./logs')])

# Evaluate
test_loss, test_accuracy = model.evaluate(test_ds)
print(f"Test accuracy: {test_accuracy:.4f}")

# Save model
model.save('mnist_model')

How This Works

  • Dockerfile: Sets up a TensorFlow environment, copies mnist.py, and installs Matplotlib.
  • mnist.py: Loads MNIST, builds a CNN, trains (~98–99% accuracy), evaluates, and saves the model, logging metrics to ./logs.
  • Execution: Docker runs mnist.py, producing a trained model and logs in your local directory.

Running the Setup

  1. Create tensorflow_project directory and add Dockerfile and mnist.py.
  2. In a terminal, navigate to tensorflow_project:
cd tensorflow_project
  1. Build the image:
docker build -t my-tensorflow-app .
  1. Run the container:
docker run -v $(pwd):/app my-tensorflow-app
  1. Expect ~1–2 minutes for training, ~98–99% accuracy, and output files (mnist_model, logs).
  2. View TensorBoard:
tensorboard --logdir ./logs

Outcome

You’ve run TensorFlow in a Docker container, trained a CNN, and saved a model, ready for deployment or further development.

Best Practices

  • Use Official Images: Start with tensorflow/tensorflow images for reliability.
  • Mount Volumes: Use -v to persist models and logs locally.
  • Optimize Images: Keep Dockerfiles lean by installing only necessary dependencies.
  • Test Locally: Run containers interactively (-it bash) to debug before automating.
  • Document: Comment your Dockerfile and scripts for clarity ([Docker for ML](/tensorflow/production/docker-for-ml)).

Troubleshooting

  • Docker Not Running: Ensure Docker Desktop is active (docker info) or restart it ([Installation Troubleshooting](/tensorflow/introduction/installation-troubleshooting)).
  • Image Pull Fails: Check internet and Docker Hub access (docker pull).
  • File Not Found: Verify mnist.py is in tensorflow_project and mounted correctly.
  • Training Errors: Ensure TensorFlow version matches (2.16.2) and data shapes are correct ([Tensor Shapes](/tensorflow/fundamentals/tensor-shapes)).
  • Help: Visit [TensorFlow Community Resources](/tensorflow/introduction/tensorflow-community-resources) or [tensorflow.org/community](https://www.tensorflow.org/community).

Next Steps

  • Explore Models: Build [RNNs](/tensorflow/advanced/recurrent-neural-networks) or [Transformers](/tensorflow/nlp/transformer-nlp) with [TensorFlow in Deep Learning](/tensorflow/introduction/tensorflow-in-deep-learning).
  • Scale Up: Deploy with [TensorFlow Serving](/tensorflow/production/tensorflow-serving) or [Cloud Integration](/tensorflow/introduction/cloud-integration).
  • Build Projects: Create [Stock Price Prediction](/tensorflow/projects/stock-price-prediction) or [TensorFlow Portfolio](/tensorflow/projects/tensorflow-portfolio).
  • Learn More: Earn [TensorFlow Certifications](/tensorflow/introduction/tensorflow-certifications).

Conclusion

Running TensorFlow with Docker simplifies machine learning by providing a consistent, portable environment. By following these steps—installing Docker, pulling TensorFlow images, creating a custom image, and training a MNIST classifier—you’ve learned to containerize TensorFlow projects, achieving ~98–99% accuracy. This approach applies to any task, from Real-Time Detection to Custom AI Solution. Explore more at tensorflow.org and check out TensorFlow Datasets or TensorFlow Documentation to keep building.