Deploy and Scale: Serving TensorFlow 2#

1. Consistency and Reliability#

When you use TensorFlow Serving or similar serving infrastructures, your model inference is more consistent compared to a custom-built, ad-hoc deployment. It provides standardized APIs that are flexible, yet production-ready, minimizing potential errors that can occur when each team rolls its own solution.

2. Performance and Scalability#

Specialized tools can help you scale easily. TensorFlow Serving is built with high performance in mind, leveraging gRPC and efficient batching under the hood. It also has built-in support for features like dynamic loading of new model versions without downtime.

3. Ease of Deployment#

Instead of bridging the gap between raw model files and a production web server, TensorFlow Serving provides an out-of-the-box structure for exposing your model. This makes it straightforward to transition from experimentation to production.

4. Flexibility#

You can serve multiple models or multiple versions of the same model in one server. This flexibility allows you to test, debug, or roll out new models in a more controlled manner.

Key Benefits#

• Versioning: TensorFlow Serving can host multiple versions, allowing your service to upgrade seamlessly.
• Batching: For heavy workloads, you can enable dynamic batching of inference requests.
• Monitoring: TensorFlow Serving exposes metrics that can be integrated with popular monitoring frameworks.

Below is a simplified conceptual table illustrating some of the core features of TensorFlow Serving:

Example: Simple Image Classification Model#

Let’s illustrate with a simple Keras-based image classification model. In this example, we will use the MNIST dataset. Although MNIST is basic, the concepts extend to more complex datasets and models.

1
import tensorflow as tf
2
from tensorflow import keras
3
import numpy as np
4

5
# Load the MNIST dataset
6
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
7
x_train, x_test = x_train / 255.0, x_test / 255.0
8

9
# Expand dimensions to match the input shape requirement
10
x_train = np.expand_dims(x_train, axis=-1)
11
x_test = np.expand_dims(x_test, axis=-1)
12

13
# Define a simple CNN model
14
model = keras.models.Sequential([
15
    keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28,28,1)),
16
    keras.layers.MaxPooling2D((2,2)),
17
    keras.layers.Flatten(),
18
    keras.layers.Dense(64, activation='relu'),
19
    keras.layers.Dense(10, activation='softmax')
20
])
21

22
model.compile(optimizer='adam',
23
              loss='sparse_categorical_crossentropy',
24
              metrics=['accuracy'])
25

26
# Train the model
27
model.fit(x_train, y_train, epochs=1, validation_split=0.1)
28

29
# Evaluate the model
30
loss, accuracy = model.evaluate(x_test, y_test)
31
print(f"Test Loss: {loss}, Test Accuracy: {accuracy}")
32

33
# Save the model in the SavedModel format
34
model.save("saved_mnist_model")

When you run this script, TensorFlow saves the model into a folder named saved_mnist_model, which includes all the necessary files for TensorFlow Serving:

• saved_model.pb (the model’s graph and weights in a protobuf format),
• a variables folder (containing checkpoint files),
• possible asset directories if your model references external files.

Prerequisites#

• TensorFlow Serving installed on your local machine, or Docker installed to run the official TensorFlow Serving Docker image.

1
# Add the TensorFlow Serving package repository
2
sudo apt-get update
3
sudo apt-get install -y gnupg
4
curl https://storage.googleapis.com/tensorflow-serving-apt/tensorflow-serving.release.pub.gpg | sudo apt-key add -
5
echo "deb [arch=amd64] http://storage.googleapis.com/tensorflow-serving-apt stable tensorflow-serving" | sudo tee /etc/apt/sources.list.d/tensorflow-serving.list
6

7
# Install TensorFlow Serving
8
sudo apt-get update
9
sudo apt-get install tensorflow-model-server

1
tensorflow_model_server \
2
  --model_base_path="/path/to/saved_mnist_model" \
3
  --rest_api_port=8501 \
4
  --model_name=mnist

Verifying the Service#

Open another terminal (or use a tool like curl or Postman) to send an inference request:

1
curl -X POST \
2
  -d '{"instances": [[[ [0.0], [0.0], ... ]]]}' \
3
  http://localhost:8501/v1/models/mnist:predict

Of course, you’ll want to replace the JSON with a properly structured MNIST image. If everything is configured correctly, you will receive a JSON response containing the predicted class probabilities.

Why Use Docker?#

Running ML models inside containers brings several advantages:

• Consistency: Environment discrepancies vanish because the container locks all dependencies.
• Portability: Containers run equally well on local machines, on-premises servers, or in the cloud.
• Scalability: Orchestrators like Kubernetes can easily scale containers horizontally.

Building Your Docker Image#

We can leverage the official TensorFlow Serving Docker image or build our own Docker container. Let’s use the official image to serve the model:

1. Pull the official TF Serving image:

   1
   docker pull tensorflow/serving
   

2. Put your saved model in a local directory (for example, /models/mnist).

3. Run a container mounting the local model directory:

   1
   docker run -p 8501:8501 \
   2
     --mount type=bind,\
   3
     source=/models/mnist,\
   4
     target=/models/mnist \
   5
     -e MODEL_BASE_PATH=/models/mnist \
   6
     -e MODEL_NAME=mnist \
   7
     tensorflow/serving
   

This mounts your host directory (/models/mnist) into the container’s file system at /models/mnist, sets environment variables that TensorFlow Serving will read, and exposes port 8501 for REST requests.

Basic Kubernetes Workflow#

1. Create a Deployment manifest that includes your container image and sets up desired replicas.
2. Expose the Deployment via a Service that allows external traffic to reach your containerized model server.
3. Auto-scaling (optional) using a Horizontal Pod Autoscaler (HPA) to scale the number of pods based on CPU or custom metrics.

Below is an example deployment.yaml for a model named mnist:

1
apiVersion: apps/v1
2
kind: Deployment
3
metadata:
4
  name: mnist-deployment
5
spec:
6
  replicas: 2
7
  selector:
8
    matchLabels:
9
      app: mnist
10
  template:
11
    metadata:
12
      labels:
13
        app: mnist
14
    spec:
15
      containers:
16
      - name: mnist-container
17
        image: tensorflow/serving
18
        ports:
19
        - containerPort: 8501
20
        volumeMounts:
21
        - name: model-volume
22
          mountPath: /models/mnist
23
        env:
24
        - name: MODEL_BASE_PATH
25
          value: /models/mnist
26
        - name: MODEL_NAME
27
          value: mnist
28
      volumes:
29
      - name: model-volume
30
        persistentVolumeClaim:
31
          claimName: mnist-pvc

Service Configuration#

Create a service.yaml file to expose this deployment:

1
apiVersion: v1
2
kind: Service
3
metadata:
4
  name: mnist-service
5
spec:
6
  selector:
7
    app: mnist
8
  ports:
9
  - protocol: TCP
10
    port: 80
11
    targetPort: 8501
12
  type: LoadBalancer

With both the deployment and service in place, you can run:

1
kubectl apply -f deployment.yaml
2
kubectl apply -f service.yaml

The Kubernetes cluster will then:

1. Pull the tensorflow/serving image.
2. Spawn the specified number of pods (replicas: 2) hosting the model.
3. Expose them via a load balancer or a NodePort (depending on your environment).

Horizontal Pod Autoscaling#

If you want to scale the deployment automatically based on CPU usage, define an HPA:

1
kubectl autoscale deployment mnist-deployment \
2
  --cpu-percent=50 \
3
  --min=1 \
4
  --max=10

This tells Kubernetes to keep average CPU usage across all Pods at around 50%. If the CPU usage spikes, K8s will spin up more Pods as needed.

1. Hardware Accelerators#

Training and serving large models using CPUs alone can be slow. Consider:

• GPUs: Offer parallel computation for speed-ups in neural network inference.
• TPUs: Tensor Processing Units for even faster computations (available on Google Cloud and certain on-prem setups).

2. Model Optimization Techniques#

• Quantization: Convert model weights (and possibly activations) to lower precision (e.g., INT8) for smaller size and faster inference.
• Pruning: Remove unnecessary weights to reduce the model’s size.
• Compiler Optimization: Use XLA or TensorRT for further optimization.

3. Batching and Thread Configuration#

TensorFlow Serving can batch multiple requests together, especially if you’re using gRPC. Adjusting the concurrency settings and thread pool size can yield substantial improvements under high load.

4. Caching#

If your model predictions have some overlap in requests, you can integrate caching at the application level. Use carefully, as many ML models have unique inputs.

5. Monitoring#

TensorFlow Serving exports metrics (e.g., average latency, request count, error rates) via Prometheus endpoints. Common steps:

1. Enable Metrics: Start the server with flags that expose metrics.
2. Scrape with Prometheus: Configure Prometheus to collect these metrics from the pods.
3. Visualize: Use Grafana to set up dashboards for real-time monitoring of inference latency and throughput.

1. TFX and End-to-End Pipelines#

TensorFlow Extended (TFX) is an end-to-end platform for deploying production ML pipelines. You can integrate TFX components such as ExampleGen, Trainer, Evaluator, and Pusher to automate the entire lifecycle from data ingestion through model validation and serving. Once a model passes all the checks, TFX can automatically push the saved model to a serving instance.

2. Continuous Integration and Continuous Deployment (CI/CD)#

For a truly production-grade system, integrate your model serving with CI/CD platforms (Jenkins, GitLab CI, GitHub Actions, etc.). For instance:

• Automated Builds: When a new model is trained, automatically package it into a Docker image.
• Automated Testing: Validate the model’s performance on a hold-out dataset or real traffic replay.
• Deploy: If it meets thresholds, push the updated Docker image to your container registry, and let Kubernetes rolling upgrades do the rest.

3. Canary Releases and A/B Testing#

With versioning in TensorFlow Serving, you can run multiple model versions side by side. This makes it straightforward to try:

• Canary Releases: Gradually shift traffic from the old model to the new model, monitoring key metrics (accuracy, latency, error rates).
• A/B Testing: Split traffic between two versions. Gather user feedback or business metrics to determine the winner.

4. Edge Serving#

For use cases requiring low-latency or offline inference, you can deploy TensorFlow models on edge devices. Tools like TensorFlow Lite optimize models to run on mobile and embedded platforms. While not strictly “served�?as in centralized servers, the same lifecycle considerations apply: versioning, monitoring, and updates.