使用python多进程实现session.run (tensorflow)的并行化-用于推理

import multiprocessing
import tensorflow as tf

def worker(graph_def, input_tensors, output_tensors, input_data, output_data):
    with tf.Session(graph=tf.Graph()) as session:
        session.graph.as_default()
        tf.import_graph_def(graph_def, name='')
        inputs = {input_tensors[i].name: input_data[i] for i in range(len(input_tensors))}
        outputs = [output_tensor.name for output_tensor in output_tensors]
        result = session.run(outputs, feed_dict=inputs)
        for i in range(len(output_data)):
            output_data[i] = result[i]

graph_def = tf.GraphDef()
with tf.gfile.FastGFile('path_to_graph.pb', 'rb') as f:
    graph_def.ParseFromString(f.read())

input_tensors = [graph_def.node[i].name for i in range(len(graph_def.node)) if 'input' in graph_def.node[i].name]
output_tensors = [graph_def.node[i].name for i in range(len(graph_def.node)) if 'output' in graph_def.node[i].name]

# 定义输入数据和输出数据
input_data = [...]
output_data = [None] * len(output_tensors)

pool = multiprocessing.Pool(processes=num_processes)

for _ in range(num_processes):
    pool.apply_async(worker, args=(graph_def, input_tensors, output_tensors, input_data, output_data))

pool.close()
pool.join()

# 处理输出数据
for i in range(len(output_data)):
    print("Output {}: {}".format(i, output_data[i]))