使用Flask通过REST API服务训练过的Tensorflow模型？

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.

Answer 3

I think that one of your main concerns might be batching the requests. For example, let's say that your model is a trained CNN like VGG, Inception or similar. If you implement a regular web service with Flask, for each prediction request you receive (assuming you're running on GPU) you will do the prediction of a single image in the GPU, which can be suboptimal since you could batch similar requests, for example.

That's one of the things that TensorFlow Serving aims to offer, being able to combine requests for the same model/signature into a single batch before sending to GPU, being more efficient in the use of resources and (potentially) in throughput. You can find more information here: https://github.com/tensorflow/serving/tree/master/tensorflow_serving/batching

That said, it depends on the scenario very much. But batching of the predictions is something important to keep in mind.