A wrapper for thread safe execution of Keras model prediction when using Tensorflow as backend

thread_safe_model.py

import threading
from keras.models import Model
import keras


class ThreadSafeModel:

    def __init__(self, model: Model):
        self.model = model
        self.lock = threading.Lock()
        self.session = keras.backend.get_session()  # store session that was used when creating or loading the model

    def predict(self, x, batch_size=None, verbose=0, steps=None, callbacks=None, max_queue_size=10, workers=1,
                use_multiprocessing=False):
        with self.lock:
            with self.session.graph.as_default():
                keras.backend.set_session(self.session)
                result = self.model.predict(
                    x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)
        return result