Keras + Tensorflow and Multiprocessing in Python

I'm using Keras with Tensorflow as backend.

I am trying to save a model in my main process and then load/run (i.e. call model.predict) within another process.

I'm currently just trying the naive approach from the docs to save/load the model: https://keras.io/getting-started/faq/#how-can-i-save-a-keras-model.
So basically:

  1. model.save() in main process
  2. model = load_model() in child process
  3. model.predict() in child process

However, it simply hangs on the load_model call.

Searching around I've discovered this potentially related answer suggesting that Keras can only be utilized in one process: using multiprocessing with theano but am unsure if this is true (can't seem to find much on this).

Is there a way to accomplish my goal? A high level description or short example is greatly appreciated.

Note: I've attempted approaches along the lines of passing a graph to the process but failed since it seems tensorflow graphs aren't pickable (related SO post for that here: Tensorflow: Passing a session to a python multiprocess). If there is indeed a way to pass the tensorflow graph/model to the child process then I am open to that as well.

Thanks!


From my experience - the problem lies in loading Keras to one process and then spawning a new process when the keras has been loaded to your main environment. But for some applications (like e.g. training a mixture of Kerasmodels) it's simply better to have all of this things in one process. So what I advise is the following (a little bit cumbersome - but working for me) approach:

  1. DO NOT LOAD KERAS TO YOUR MAIN ENVIRONMENT. If you want to load Keras / Theano / TensorFlow do it only in the function environment. E.g. don't do this:

    import keras
    
    def training_function(...):
        ...
    

    but do the following:

    def training_function(...):
        import keras
        ...
    
  2. Run work connected with each model in a separate process: I'm usually creating workers which are making the job (like e.g. training, tuning, scoring) and I'm running them in separate processes. What is nice about it that whole memory used by this process is completely freed when your process is done. This helps you with loads of memory problems which you usually come across when you are using multiprocessing or even running multiple models in one process. So this looks e.g. like this:

    def _training_worker(train_params):
        import keras
        model = obtain_model(train_params)
        model.fit(train_params)
        send_message_to_main_process(...)
    
    def train_new_model(train_params):
        training_process = multiprocessing.Process(target=_training_worker, args = train_params)
        training_process.start()
        get_message_from_training_process(...)
        training_process.join()
    

Different approach is simply preparing different scripts for different model actions. But this may cause memory errors especially when your models are memory consuming. NOTE that due to this reason it's better to make your execution strictly sequential.


I created one simple example to show how to run Keras model in multiple processes with multiple gpus. Hope this sample could help you. https://github.com/yuanyuanli85/Keras-Multiple-Process-Prediction


I created a decorator that fixed my code.

from multiprocessing import Pipe, Process

def child_process(func):
    """Makes the function run as a separate process."""
    def wrapper(*args, **kwargs):
        def worker(conn, func, args, kwargs):
            conn.send(func(*args, **kwargs))
            conn.close()
        parent_conn, child_conn = Pipe()
        p = Process(target=worker, args=(child_conn, func, args, kwargs))
        p.start()
        ret = parent_conn.recv()
        p.join()
        return ret
return wrapper

@child_process
def keras_stuff():
    """ Keras stuff here"""