Tensorflow load model using multiprocessing python

Viewed 471

I have to load tensorflow trained models in parallel to minimize the loading time. For this, I'm using multiprocessing module of Python. After the completion of loading a model, it returns a function and I'm saving this function in a shared dictionary obtained by multiprocessing.Manager(). While adding returned function in the dictionary, I getting this error:

AttributeError: Can't pickle local object 'Loader._recreate_base_user_object.._UserObject'

Importing required modules and creating global dictionaries:

import tensorflow.compat.v2 as tf_v2
from multiprocessing import Process, Manager

manager = Manager()

detection_functions_dict = manager.dict()

Function to start new Process for every model to load:

def load_models_async(inference_graph_array):
    for inference_graph in inference_graph_array:
        process = Process(target=load_model, args=(inference_graph,))
        process.start()

Function to load model:

def load_model(inference_graph):
    try:
        detection_function = tf_v2.saved_model.load(inference_graph.model_path)

        # Getting error at below line
        detection_functions_dict[inference_graph.identifier] = detection_function

Calling this from main thread:

if __name__ == '__main__':
    port = int(os.environ.get('APP_PORT', 5004))
    host = os.environ.get('APP_HOST', '0.0.0.0')

    # inference_graph is a class having identifier and model path variables
    load_models_async(inference_graph_array)

    app.run(host=host, port=port)

How can I resolve this error?

0 Answers
Related