TensorFlow
RESTful API
model deployment
machine learning
API integration

Example for Deploying a Tensorflow Model via a RESTful API

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Deploying a TensorFlow model behind a REST API lets other applications request predictions over HTTP without embedding TensorFlow directly. A simple way to do this is to load a saved model in a Python web service and expose a /predict endpoint. For learning and internal tools, Flask is enough. For larger production systems, the same concepts still apply, even if you later move to TensorFlow Serving or another inference platform.

Save or Load a Trained Model

Assume you already have a saved TensorFlow model:

python
import tensorflow as tf

model = tf.keras.models.load_model("saved_model/my_model")

The model should accept a tensor shape that your API can reproduce from JSON input. If your service receives feature vectors, the endpoint should turn incoming values into a batch-shaped tensor before calling the model.

Build a Minimal Flask API

Here is a small Flask app that loads the model once at startup and serves predictions.

python
1from flask import Flask, jsonify, request
2import tensorflow as tf
3
4app = Flask(__name__)
5model = tf.keras.models.load_model("saved_model/my_model")
6
7@app.post("/predict")
8def predict():
9    payload = request.get_json()
10
11    if "instances" not in payload:
12        return jsonify({"error": "Missing 'instances' field"}), 400
13
14    x = tf.convert_to_tensor(payload["instances"], dtype=tf.float32)
15    predictions = model(x, training=False)
16
17    return jsonify({
18        "predictions": predictions.numpy().tolist()
19    })
20
21if __name__ == "__main__":
22    app.run(host="0.0.0.0", port=5000)

This is the essential flow:

  1. parse JSON
  2. validate the expected field
  3. convert input to a tensor
  4. run inference
  5. return JSON

Call the Endpoint

Once the server is running, send a request like:

bash
1curl -X POST http://localhost:5000/predict \
2  -H "Content-Type: application/json" \
3  -d '{
4    "instances": [
5      [5.1, 3.5, 1.4, 0.2]
6    ]
7  }'

A typical response looks like:

json
{
  "predictions": [[0.97, 0.02, 0.01]]
}

The exact shape depends on your model. Classification models often return class probabilities or logits, while regression models may return a single number per example.

Prepare Inputs Consistently

The API is only as good as its preprocessing. If the model was trained on normalized features, tokenized text, or resized images, the same steps must happen inside the API service.

For example, image models often need:

  • decode image bytes
  • resize to the training resolution
  • normalize pixel values
  • add a batch dimension

If you skip or change preprocessing, the model may produce confident but wrong predictions.

Return Useful Errors

Do not let malformed requests fail with vague stack traces. Validate input shape and required fields early.

python
if not isinstance(payload["instances"], list):
    return jsonify({"error": "'instances' must be a list"}), 400

Simple validation makes the API much easier to integrate and debug.

Production Considerations

A Flask example is good for learning, prototypes, and some internal services, but production inference often needs more:

  • structured logging
  • health endpoints
  • request timeouts
  • batching or concurrency control
  • model versioning
  • authentication and HTTPS

For heavier traffic, TensorFlow Serving or another dedicated inference system is usually a better long-term choice than hand-building every deployment concern into a Flask app.

Common Pitfalls

The most common mistake is loading the model inside the request handler. That makes every request slow and wastes resources. Load the model once at process startup.

Another issue is forgetting the batch dimension. Many models expect shape [batch, features], even for one example.

Developers also often overlook preprocessing consistency. The deployed endpoint must mirror training-time input preparation.

Finally, remember that .numpy() works naturally in eager mode, but the returned values still need conversion to plain Python types such as lists before JSON serialization.

Summary

  • A REST API around TensorFlow usually loads a saved model once and exposes a /predict endpoint.
  • Flask is enough for a simple working example.
  • Convert incoming JSON to tensors, run inference, then return JSON-friendly output.
  • Keep preprocessing consistent with training.
  • For production-scale serving, consider a dedicated inference platform after validating the basic API flow.

Course illustration
Course illustration

All Rights Reserved.