Serve trained Tensorflow model with REST API using Flask?

TensorFlow

Flask

REST API

machine learning deployment

model serving

Serve trained Tensorflow model with REST API using Flask?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

TensorFlow & Keras Machine Learning DevOps & Deployment Networking & APIs Python

Introduction

Serving a trained TensorFlow model through a Flask REST API involves loading the saved model at startup, creating an endpoint that accepts input data (usually JSON), preprocessing the input into the format the model expects, running inference, and returning the prediction as a JSON response. This approach is suitable for prototyping and low-traffic deployments. For production use, consider TensorFlow Serving or FastAPI with async support.

Saving a Trained Model

First, save your trained model in the SavedModel or HDF5 format:

python

1import tensorflow as tf
2
3# After training
4model = tf.keras.Sequential([
5    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
6    tf.keras.layers.Dropout(0.2),
7    tf.keras.layers.Dense(10, activation='softmax')
8])
9
10model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
11model.fit(x_train, y_train, epochs=5)
12
13# Save as SavedModel (recommended)
14model.save('saved_model/mnist_model')
15
16# Or save as HDF5
17model.save('mnist_model.h5')

Basic Flask API

python

1from flask import Flask, request, jsonify
2import tensorflow as tf
3import numpy as np
4
5app = Flask(__name__)
6
7# Load model once at startup (not per request)
8model = tf.keras.models.load_model('saved_model/mnist_model')
9
10@app.route('/predict', methods=['POST'])
11def predict():
12    data = request.get_json()
13
14    # Convert input to numpy array
15    input_data = np.array(data['instances'])
16
17    # Run inference
18    predictions = model.predict(input_data)
19
20    # Convert to list for JSON serialization
21    results = predictions.tolist()
22
23    return jsonify({
24        'predictions': results
25    })
26
27if __name__ == '__main__':
28    app.run(host='0.0.0.0', port=5000)

Test the endpoint:

bash

curl -X POST http://localhost:5000/predict \
  -H "Content-Type: application/json" \
  -d '{"instances": [[0.1, 0.2, 0.3, ...]]}'

Input Preprocessing

Most models require specific input shapes and normalization:

python

1@app.route('/predict', methods=['POST'])
2def predict():
3    data = request.get_json()
4
5    # Validate input
6    if 'instances' not in data:
7        return jsonify({'error': 'Missing "instances" field'}), 400
8
9    try:
10        input_data = np.array(data['instances'], dtype=np.float32)
11    except (ValueError, TypeError) as e:
12        return jsonify({'error': f'Invalid input format: {str(e)}'}), 400
13
14    # Validate shape
15    expected_shape = model.input_shape[1:]  # e.g., (784,)
16    if input_data.shape[1:] != expected_shape:
17        return jsonify({
18            'error': f'Expected shape {expected_shape}, got {input_data.shape[1:]}'
19        }), 400
20
21    # Normalize if needed (e.g., pixel values 0-255 to 0-1)
22    input_data = input_data / 255.0
23
24    predictions = model.predict(input_data)
25
26    return jsonify({
27        'predictions': predictions.tolist()
28    })

Image Classification Endpoint

For serving an image classification model:

python

1from flask import Flask, request, jsonify
2import tensorflow as tf
3import numpy as np
4from PIL import Image
5import io
6
7app = Flask(__name__)
8model = tf.keras.models.load_model('saved_model/image_classifier')
9
10CLASS_NAMES = ['cat', 'dog', 'bird', 'fish']
11
12@app.route('/classify', methods=['POST'])
13def classify():
14    if 'image' not in request.files:
15        return jsonify({'error': 'No image file provided'}), 400
16
17    file = request.files['image']
18    image = Image.open(io.BytesIO(file.read()))
19
20    # Resize to model input size
21    image = image.resize((224, 224))
22    image_array = np.array(image) / 255.0
23    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
24
25    predictions = model.predict(image_array)
26    predicted_class = CLASS_NAMES[np.argmax(predictions[0])]
27    confidence = float(np.max(predictions[0]))
28
29    return jsonify({
30        'class': predicted_class,
31        'confidence': confidence,
32        'probabilities': dict(zip(CLASS_NAMES, predictions[0].tolist()))
33    })

Test with an image file:

bash

curl -X POST http://localhost:5000/classify \
  -F "[email protected]"

Health Check and Model Info

python

1@app.route('/health', methods=['GET'])
2def health():
3    return jsonify({'status': 'healthy'})
4
5@app.route('/model/info', methods=['GET'])
6def model_info():
7    return jsonify({
8        'input_shape': str(model.input_shape),
9        'output_shape': str(model.output_shape),
10        'num_parameters': int(model.count_params())
11    })

Production Considerations

Use Gunicorn Instead of Flask Dev Server

bash

1pip install gunicorn
2
3# Run with 4 worker processes
4gunicorn --workers 4 --bind 0.0.0.0:5000 app:app
5
6# With timeout for slow predictions
7gunicorn --workers 4 --timeout 120 --bind 0.0.0.0:5000 app:app

Thread Safety with TensorFlow

TensorFlow is not fully thread-safe. Use a lock for predictions in multi-threaded environments:

python

1import threading
2
3model_lock = threading.Lock()
4
5@app.route('/predict', methods=['POST'])
6def predict():
7    data = request.get_json()
8    input_data = np.array(data['instances'], dtype=np.float32)
9
10    with model_lock:
11        predictions = model.predict(input_data)
12
13    return jsonify({'predictions': predictions.tolist()})

Batch Predictions

python

1@app.route('/batch_predict', methods=['POST'])
2def batch_predict():
3    data = request.get_json()
4    instances = np.array(data['instances'], dtype=np.float32)
5
6    # Process in batches of 32
7    batch_size = 32
8    all_predictions = []
9
10    for i in range(0, len(instances), batch_size):
11        batch = instances[i:i + batch_size]
12        preds = model.predict(batch, batch_size=batch_size)
13        all_predictions.extend(preds.tolist())
14
15    return jsonify({'predictions': all_predictions})

Docker Deployment

dockerfile

1FROM python:3.10-slim
2
3WORKDIR /app
4
5COPY requirements.txt .
6RUN pip install --no-cache-dir -r requirements.txt
7
8COPY saved_model/ ./saved_model/
9COPY app.py .
10
11EXPOSE 5000
12
13CMD ["gunicorn", "--workers", "2", "--bind", "0.0.0.0:5000", "app:app"]

1# requirements.txt
2flask==3.0.0
3tensorflow==2.15.0
4gunicorn==21.2.0
5numpy==1.26.0
6Pillow==10.1.0

Common Pitfalls

Loading the model inside the request handler: tf.keras.models.load_model() takes seconds to minutes for large models. Load the model once at module level or in an app.before_first_request handler, not inside the endpoint function.
Using Flask's development server in production: app.run() uses a single-threaded server that cannot handle concurrent requests. Use Gunicorn or uWSGI in production with multiple workers.
NumPy arrays are not JSON serializable: model.predict() returns a NumPy array. Call .tolist() before passing it to jsonify(), or the response will fail with a TypeError.
Forgetting the batch dimension: Most TensorFlow models expect input shape (batch_size, ...). A single input of shape (784,) must be reshaped to (1, 784) with np.expand_dims(input_data, axis=0).
GPU memory with multiple workers: Each Gunicorn worker loads a separate copy of the model into GPU memory. With 4 workers and a large model, you may run out of GPU memory. Use tf.config.set_memory_growth(gpu, True) or limit workers to 1-2 for GPU inference.

Summary

Save models with model.save() and load at Flask app startup with tf.keras.models.load_model()
Create a /predict endpoint that accepts JSON input, runs model.predict(), and returns JSON
Validate input shape and data types before running inference
Use Gunicorn with multiple workers for production deployment
Add a threading lock around model.predict() for thread safety
Use Docker for reproducible deployments with pinned dependency versions

System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Serve trained Tensorflow model with REST API using Flask?

ML System Design practice on Codemia

Introduction

Saving a Trained Model

Basic Flask API

Input Preprocessing

Image Classification Endpoint

Health Check and Model Info

Production Considerations

Use Gunicorn Instead of Flask Dev Server

Thread Safety with TensorFlow

Batch Predictions

Docker Deployment

Common Pitfalls

Summary

Related reading

System Design Fundamentals

Track what you have practised

ML System Design practice on Codemia