Serve trained Tensorflow model with REST API using Flask?

Question

I have got a trained Tensorflow model and I want to serve the prediction method with REST API. What I can think of is to use Flask to build a simple REST API that receive JSON as input and then call the predict method in Tensorflow and then return the predicted result to the client side.

I would like to know is there any concern to do it this way especially in production environment?

Many thanks!

Answer 1

The first concern which comes into my mind is the performance.

TensorFlow team seems to have worked out server/client usage. You may want to look into tensorflow serving. As a default, it uses gRPC for communication protocol.

Answer 2

We use Flask + TensorFlow serving at work. Our setup might not be the most optimal way to serve models, but it gets the job done and it works fine for us so far.

The setup is the following:

1. Because tfserving takes forever to build, we built a docker image (not GPU support or anything, but it works for just serving a model and it's faster and better than serving it directly from within a huge Python/Flask monolite). The model server image can be found here: https://hub.docker.com/r/epigramai/model-server/
2. Then Flask is used to setup an API. In order to send requests to the model server we need a grcp prediction client, so we built one in Python that we can import directly into the flask API, https://github.com/epigramai/tfserving_predict_client/.

The good thing here is that the model is not served by the Flask API application. The docker image model server can easily be replaced with a model server running on a GPU compiled for the machines hardware instead of the docker container.