TensorFlow中双PB模型并行推理及结果融合实现问询

当然可以搞定！这种整合不仅能省去开多个终端的麻烦，还能更高效地利用你的硬件资源。我给你梳理下具体的实现思路和代码示例，你可以照着改：

核心思路 #

本质就是三步：把两个PB模型都加载到同一个脚本里，用并行方式让它们同时处理同一帧图像，最后把两个输出按需求合并。这里的并行可以用多线程（适合GPU场景）或者多进程（适合CPU场景），TensorFlow对这两种方式都支持。

具体实现步骤 #

1. 封装通用的PB模型加载函数 #

因为两个模型都是PB格式，我们可以写一个复用的加载函数，避免重复代码。关键是要确认你模型的输入输出张量名称，这个可以用TensorFlow的工具或者在加载后打印所有张量名称来获取。

2. 并行执行模型推理 #

用Python的concurrent.futures.ThreadPoolExecutor实现并行，这样两个模型可以同时处理同一帧图像，不用等一个跑完再跑另一个。如果你的模型以CPU推理为主，也可以换成ProcessPoolExecutor（多进程），避开GIL的性能限制。

3. 合并输出结果 #

拿到两个模型的输出后，根据你的业务需求做叠加——比如简单的加权求和、取平均值，或者更复杂的融合逻辑，完全按需调整。

完整代码示例 #

import tensorflow as tf
import cv2
from concurrent.futures import ThreadPoolExecutor

# 封装PB模型加载函数，兼容TF1格式的PB文件
def load_pb_model(pb_path):
    graph = tf.compat.v1.Graph()
    with graph.as_default():
        graph_def = tf.compat.v1.GraphDef()
        with tf.io.gfile.GFile(pb_path, 'rb') as fid:
            serialized_graph = fid.read()
            graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(graph_def, name='')
    
    # 配置GPU显存按需分配（使用GPU时打开，避免显存占满）
    config = tf.compat.v1.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.compat.v1.Session(graph=graph, config=config)
    
    # 替换成你自己模型的实际输入输出张量名称！
    input_tensor = graph.get_tensor_by_name('input:0')  # 示例名称，需修改
    output_tensor = graph.get_tensor_by_name('output:0') # 示例名称，需修改
    
    return sess, input_tensor, output_tensor

# 推理函数，包含图像预处理（必须和原脚本一致）
def run_inference(sess, input_tensor, output_tensor, image):
    # 替换成你原有脚本里的图像预处理逻辑
    # 示例：resize到模型要求尺寸、归一化、增加batch维度
    preprocessed = cv2.resize(image, (224, 224))
    preprocessed = preprocessed / 255.0
    preprocessed = preprocessed[tf.newaxis, ...]  # 增加batch维度
    
    # 执行推理
    output = sess.run(output_tensor, feed_dict={input_tensor: preprocessed})
    return output

# 输出后处理（按需调整）
def postprocess_output(final_output):
    # 示例：如果是图像输出，保存并显示
    if len(final_output.shape) == 4:  # 去除batch维度
        final_output = final_output[0]
    final_output = (final_output * 255).astype('uint8')
    cv2.imwrite('final_result.jpg', final_output)
    cv2.imshow('Final Result', final_output)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def main():
    # 加载两个模型
    sess_a, input_a, output_a = load_pb_model('A.pb')
    sess_b, input_b, output_b = load_pb_model('B.pb')

    # 读取图像（如果是实时帧，替换成摄像头读取循环即可）
    image = cv2.imread('test_frame.jpg')

    # 并行执行两个模型的推理
    with ThreadPoolExecutor(max_workers=2) as executor:
        future_a = executor.submit(run_inference, sess_a, input_a, output_a, image)
        future_b = executor.submit(run_inference, sess_b, input_b, output_b, image)
        
        # 获取推理结果
        res_a = future_a.result()
        res_b = future_b.result()

    # 合并输出（示例用加权平均，权重可自定义）
    final_output = 0.6 * res_a + 0.4 * res_b  # 比如给A模型更高权重

    # 处理并展示结果
    postprocess_output(final_output)

    # 关闭会话释放资源
    sess_a.close()
    sess_b.close()

if __name__ == '__main__':
    main()

关键注意事项 #

• 张量名称必须正确：一定要替换代码里的input:0和output:0为你模型的实际张量名称，不然会报错。可以用tf.compat.v1.get_default_graph().get_operations()列出所有张量名称。
• 预处理逻辑要和原脚本一致：两个模型的图像预处理（resize尺寸、归一化范围、通道顺序等）必须和你原来单独运行的脚本完全相同，否则输出结果会不符合预期。
• GPU资源管理：使用GPU时，开启显存按需分配可以避免占满显存。如果两个模型体积都很大，也可以考虑动态切换图，但并行方式的效率更高。
• 实时帧优化：如果处理摄像头实时流，建议把并行推理逻辑放到帧循环里，并用队列或异步方式避免帧堆积，保证实时性。

内容的提问来源于stack exchange，提问作者boss