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:
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
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:
curl -X POST http://localhost:5000/predict \
-H "Content-Type: application/json" \
-d '{"instances": [[0.1, 0.2, 0.3, ...]]}'
Most models require specific input shapes and normalization:
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:
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:
Health Check and Model Info
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
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:
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
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
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
