TensorFlow: Deploy multiple models with Flask

Full code : https://github.com/xiching/tf_flask/

While developing and testing machine learning models, one might want to deploy multiple models from a single web server, say, Flask/some_server.

Each model has its own Graph and Session which might have been trained in different notebooks/modules. So there could not have been the danger of mixing Sessions and/or Graphs. Let’s assume this and go ahead that the models were trained in different “context” or more naturally in different notebooks.

How to deploy these models from the same web app?

1. Serve different models from the same web app/server.
2. No two models should interfere with each other’s Graph and/or Session.
3. For the purpose of serving, unfortunately, these models have to be ‘in-memory’.

How to achieve a reasonable solution?

In this post, we will define
1) Two different models/graphs with trainable variables in them (so this means that you can do backprop on it on your own), and save these models,
2) Restore both of them such that the Sessions and/or Graphs of these two models do not interfere with each other.

Define and save models
1. First Model

In above script, we defined a model, trained it (or not), and saved it for later use.

2. Second Model

This is our second model, defined, trained/not, and saved.

Now that these two models have been defined and saved, let’s define wrapper classes which will keep the Graphs and Sessions with the class object. As shown below:

The above module, model_wrappers.py creates wrappers for our previously trained models model_a and model_b. The classes from model_wrappers can now be directly imported from flask app for serving. Let’s see how exactly.

That’s it! ‘a’ is an instance of the first model and ‘b’ is an instance of the second model. This almost looks like sklearn predict function. No worries about Sessions and Graphs interfering with one another. The model_wrappers module hides all the complexity.

If anyone has any better and prettier version of doing the same thing, please enlighten me!!

Written by

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store