Thread-safe graph usage in Tensorflow

Viewed 1677

I have a flask app that first of all loads a keras model and then executes the prediction function. According to this answer I save the tensorflow graph in a global variable and use the same graph for the prediction function.

def load_model():
    load_model_from_file()
    global graph
    graph = tf.get_default_graph()

def predict():
    with graph.as_default():
        tagger = Tagger(self.model, preprocessor=self.p)
        return tagger.analyze(words)

@app.route('/predict', methods=['GET'])
def load_and_predict():
    load_model()
    predict()

However, this leads to a problem whenever multiple requests are the send to the server. How can I make this code thread-safe, or to be more specific, how do I use tensorflow graphs correctly in a multi-threaded environment?

2 Answers

usually, you should use a session while working with threads in tensorflow.

intra_parallel_thread_tf = 1
inter_parallel_thread_tf = 1

session_conf = tf.ConfigProto(intra_op_parallelism_threads=intra_parallel_thread_tf,
                          inter_op_parallelism_threads=inter_parallel_thread_tf)

tf.Session(graph=tf.get_default_graph(), config=session_conf)
GRAPH = tf.get_default_graph()

But this is pretty general. It also depends on the error you get.

You can make it synchronized with a lock.

import threading    
lock = threading.Lock()

def load_and_predict():
     with lock:
        load_model()
        predict()
Related