TensorFlow多进程预测卡滞于session.run问题求助

这种卡在session.run的情况我之前也碰到过，核心原因是TensorFlow的计算图（Graph）和会话（Session）是进程内绑定的资源——主进程预加载的模型没法直接在子进程里复用，进程间内存是隔离的，而且TensorFlow的模型、会话这类对象大多不能被序列化跨进程传递，哪怕你用Manager或者直接传参，子进程里拿到的模型其实是无效的，最终导致会话一直阻塞等待资源。

下面给你两个最有效的解决方案，你可以根据自己的场景选择：

方案1：每个子进程独立加载模型（最推荐） #

不要在主进程提前加载所有模型，而是让每个子进程根据ID自己加载对应的模型文件。这样每个进程拥有完全独立的Graph和Session，彻底避免资源冲突。

修改后的代码示例：

def predict(id, args):
    message, configs, vocabs, emoji_dict, result_list = args
    
    # 子进程内单独加载对应ID的模型
    # 替换成你实际的模型加载逻辑，比如从文件读取
    model = load_your_model(f"model_{id}.h5", configs)
    
    # TF1.x需要创建会话；TF2.x可跳过这步，直接用eager模式预测
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        model.set_session(sess)
        
        # 执行你的预测逻辑
        pred_result = model.predict(message)
        result_list.append(pred_result)

# 主进程代码
with Manager() as manager:
    first_level_test_features = manager.list()
    procs = []
    for id in range(4):
        # 不用传预加载的models，只传必要配置和结果容器
        p = Process(target=predict, args=(id, (message, configs, vocabs, emoji_dict, first_level_test_features)))
        procs.append(p)
        p.start()
    
    for p in procs:
        p.join()

方案2：子进程创建独立计算图（若需复用主进程权重） #

如果不想重复加载完整模型文件，也可以在主进程保存每个模型的权重，让子进程重新构建模型结构并加载权重，同时创建独立的计算图：

def predict(id, args):
    message, weight_path, configs, vocabs, emoji_dict, result_list = args
    
    # 给子进程创建独立的计算图，避免和主进程冲突
    with tf.Graph().as_default() as local_graph:
        with tf.Session(graph=local_graph) as sess:
            # 根据配置重新构建模型结构
            model = build_your_model(configs)
            # 加载预保存的权重文件
            model.load_weights(weight_path)
            
            # 执行预测
            pred_result = model.predict(message)
            result_list.append(pred_result)

# 主进程代码：预先准备好每个模型的权重路径
with Manager() as manager:
    first_level_test_features = manager.list()
    model_weight_paths = [
        "model_0_weights.h5",
        "model_1_weights.h5",
        "model_2_weights.h5",
        "model_3_weights.h5"
    ]
    procs = []
    for id in range(4):
        p = Process(target=predict, args=(id, (message, model_weight_paths[id], configs, vocabs, emoji_dict, first_level_test_features)))
        procs.append(p)
        p.start()
    
    for p in procs:
        p.join()

关键注意事项 #

• 别用Manager传递模型对象：Manager仅能处理可序列化的简单对象，TensorFlow的模型、Graph、Session都无法被序列化，强行传递只会得到无效对象，引发阻塞。
• TF2.x适配：如果使用TF2.x的eager模式，也要确保每个进程独立构建模型，不要共享主进程的模型实例，否则同样会出现资源冲突。
• 进程数建议：进程数不要超过CPU核心数，过多进程会因上下文切换降低整体效率。

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