TensorFlow Serving with REST API

TensorFlow Serving is a production-grade server that loads SavedModels and exposes them as REST and gRPC endpoints. Any application — a web backend, a mobile app server, a data pipeline — can send HTTP requests and receive predictions without running Python or knowing anything about TensorFlow. TF Serving handles model versioning, hardware acceleration, request batching, and zero-downtime model swaps automatically.

Why Use TF Serving Instead of Flask or FastAPI

Flask/FastAPI approach:
  You write:
    - Python prediction code
    - Request parsing
    - Response formatting
    - Error handling
    - Batching logic
    - GPU management
  Problems: Not optimized for ML, hard to scale, manual batching

TF Serving approach:
  You provide:
    - SavedModel directory
  TF Serving handles everything else:
    - REST and gRPC endpoints automatically
    - Request batching (combines simultaneous requests)
    - GPU inference acceleration
    - Multiple model versions
    - Health checks
    - Monitoring metrics

Step 1 — Export the Model in TF Serving Layout

import tensorflow as tf
import os

# TF Serving expects: model_name/version_number/
model_name = 'image_classifier'
version    = 1
export_dir = f'serving_models/{model_name}/{version}'

model.save(export_dir)
print(f"Model saved to: {export_dir}")

# Resulting structure:
# serving_models/
# └── image_classifier/
#     └── 1/
#         ├── saved_model.pb
#         └── variables/

Step 2 — Start TF Serving With Docker

# Pull the TF Serving Docker image
docker pull tensorflow/serving

# Start the server
docker run -d --name tf_serving \
  -p 8501:8501 \
  --mount type=bind,\
source=$(pwd)/serving_models/image_classifier,\
target=/models/image_classifier \
  -e MODEL_NAME=image_classifier \
  tensorflow/serving

# The server is now running at:
# REST API:  http://localhost:8501/v1/models/image_classifier
# gRPC API:  localhost:8500

Step 3 — Check Server Health

import requests

# Check model status
response = requests.get(
    'http://localhost:8501/v1/models/image_classifier'
)
print(response.json())

# Expected output:
# {
#   "model_version_status": [{
#     "version": "1",
#     "state": "AVAILABLE",
#     "status": {"error_code": "OK"}
#   }]
# }

Step 4 — Send a Prediction Request

import requests
import json
import numpy as np

# Prepare input data
sample_image = np.random.random((1, 224, 224, 3)).astype('float32')

# TF Serving REST format: {"instances": [list_of_inputs]}
payload = json.dumps({
    'instances': sample_image.tolist()
})

# Send POST request to the predict endpoint
response = requests.post(
    'http://localhost:8501/v1/models/image_classifier:predict',
    data=payload,
    headers={'Content-Type': 'application/json'}
)

# Parse the response
result = response.json()
predictions = result['predictions']
print(f"Prediction shape: {np.array(predictions).shape}")
# (1, 10) — 10 class probabilities
predicted_class = np.argmax(predictions[0])
print(f"Predicted class: {predicted_class}")

REST API Endpoints

Base URL: http://localhost:8501/v1/models/{model_name}

Endpoint                                  Method  Purpose
──────────────────────────────────────────────────────────────────────────
/v1/models/{name}                         GET     Check model status
/v1/models/{name}/versions/{ver}          GET     Check specific version
/v1/models/{name}:predict                 POST    Get predictions (latest version)
/v1/models/{name}/versions/{ver}:predict  POST    Get predictions (specific version)
/v1/models/{name}/metadata                GET     Get input/output signature
──────────────────────────────────────────────────────────────────────────

Inspecting the Model Signature

response = requests.get(
    'http://localhost:8501/v1/models/image_classifier/metadata'
)
metadata = response.json()
print(json.dumps(metadata, indent=2))

# Shows:
# - Input tensor name, shape, dtype
# - Output tensor name, shape, dtype
# Use these to format your request payload correctly

Batching Requests for Efficiency

# Send multiple samples in one request
batch_images = np.random.random((32, 224, 224, 3)).astype('float32')

payload = json.dumps({
    'instances': batch_images.tolist()
})

response = requests.post(
    'http://localhost:8501/v1/models/image_classifier:predict',
    data=payload,
    headers={'Content-Type': 'application/json'}
)

predictions = response.json()['predictions']
print(f"Batch predictions shape: {np.array(predictions).shape}")  # (32, 10)

Model Versioning: Zero-Downtime Updates

# Deploy a new version alongside the existing one
model_v2.save('serving_models/image_classifier/2/')

# TF Serving automatically:
# 1. Detects the new version directory
# 2. Loads version 2 into memory
# 3. Routes new requests to version 2
# 4. Unloads version 1 after a grace period

# Directory now:
# serving_models/image_classifier/1/  ← old version (being drained)
# serving_models/image_classifier/2/  ← new version (serving)

# Rollback is instant — just delete version 2 folder
# TF Serving switches back to version 1 automatically

Enabling TF Serving Batching Config

# Create a batching config file
batching_config = """
max_batch_size { value: 64 }
batch_timeout_micros { value: 5000 }
num_batch_threads { value: 4 }
max_enqueued_batches { value: 100 }
"""

with open('batching_config.txt', 'w') as f:
    f.write(batching_config)

# Start server with batching enabled
docker run -d -p 8501:8501 \
  --mount type=bind,source=$(pwd)/serving_models/image_classifier,target=/models/image_classifier \
  --mount type=bind,source=$(pwd)/batching_config.txt,target=/batching_config.txt \
  -e MODEL_NAME=image_classifier \
  tensorflow/serving \
  --enable_batching=true \
  --batching_parameters_file=/batching_config.txt

GPU Support in TF Serving

# Use the GPU-enabled TF Serving image
docker pull tensorflow/serving:latest-gpu

docker run -d --gpus all -p 8501:8501 \
  --mount type=bind,source=$(pwd)/serving_models/image_classifier,target=/models/image_classifier \
  -e MODEL_NAME=image_classifier \
  tensorflow/serving:latest-gpu

Production Deployment Checklist

Before deploying to production:
  ☐ Verify model accuracy on held-out test set
  ☐ Test the REST endpoint with representative inputs
  ☐ Test with edge cases (empty input, wrong shape, wrong dtype)
  ☐ Set up health check monitoring
  ☐ Configure request batching for throughput
  ☐ Enable GPU if latency is critical
  ☐ Set up model version rollback procedure
  ☐ Configure logging for prediction auditing

TensorFlow Serving transforms a trained model into a scalable, maintainable production API with minimal setup. The combination of SavedModel export, Docker deployment, and REST endpoint gives you everything needed to serve predictions to real applications at any scale. The next topic covers TensorFlow.js — the tool for running models directly in web browsers without any server.

Previous lesson

Back to course

Next lesson