如何让多ProcessPoolExecutor进程共享单个TensorFlow模型类实例？

需要在ProcessPoolExecutor多进程环境下共享单个TensorFlow模型实例（避免4GB模型重复加载占用单GPU），同时用多进程处理计算密集的小批量数据生成，解决ThreadPoolExecutorCPU利用率低的问题。当前代码中全局模型实例会在每个子进程中重复初始化，无法实现共享。

ProcessPoolExecutor的子进程会复制父进程内存空间（Unix下fork）或重新导入模块执行代码（Windows/macOS下spawn），全局变量会被每个子进程重新实例化；且跨进程的Lock无法同步，因为每个进程的锁是独立实例。

将模型单独放在一个进程中作为服务，用队列传递推理请求和结果，多进程仅负责数据生成，既保证模型单实例，又充分利用多核心。

代码实现 #

import concurrent.futures as cf
import numpy as np
import random
from multiprocessing import Process, Queue
import tensorflow as tf  # 实际使用时导入

NUM_WORKERS = 7  # 预留1个进程给模型服务
tInfo_M = {
    'specificmodelname': []
}
CYCAVAIL = np.arange(100)
STATES = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

class DDQN:
    def __init__(self, saved_model_path):
        print(f"仅初始化一次模型：{saved_model_path}")
        # 实际加载TensorFlow模型
        self.model = tf.keras.models.load_model(saved_model_path)

    def get_qs(self, instances):
        # 替换为实际推理逻辑
        return self.model.predict(instances, verbose=0)

def model_service(request_queue, result_queue, model_path):
    """单独进程运行的模型服务，处理推理请求"""
    model = DDQN(model_path)
    while True:
        task_id, batch_data = request_queue.get()
        if task_id is None:  # 接收结束信号
            break
        q_vals = model.get_qs(batch_data)
        result_queue.put((task_id, q_vals))

def generate_and_request(task_id, request_queue, result_queue):
    """多进程执行的任务：生成数据并请求推理"""
    an_eval = []
    for _ in np.arange(30_000):
        # 生成小批量数据（此处用STATES模拟，实际替换为你的数据生成逻辑）
        batch_data = np.random.choice(STATES, size=(32,))  # 示例批量大小
        # 发送推理请求
        request_queue.put((task_id, batch_data))
        # 获取推理结果
        _, q_val = result_queue.get()
        an_eval.append(q_val)
    return an_eval

def main():
    model_path = 'specificmodelname'
    # 创建请求队列和结果队列
    request_queue = Queue(maxsize=10)  # 设置队列大小防止内存溢出
    result_queue = Queue()

    # 启动模型服务进程
    model_process = Process(
        target=model_service,
        args=(request_queue, result_queue, model_path)
    )
    model_process.start()

    # 启动多进程数据生成任务
    with cf.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
        tasks = [
            executor.submit(generate_and_request, idx, request_queue, result_queue)
            for idx, _ in enumerate(CYCAVAIL)
        ]

        # 汇总结果
        for future in cf.as_completed(tasks):
            tInfo_M['specificmodelname'].extend(future.result())

    # 发送结束信号，关闭模型服务进程
    request_queue.put((None, None))
    model_process.join()

    print(f"总结果数：{len(tInfo_M['specificmodelname'])}")

if __name__ == '__main__':
    main()

方案优势 #

1. 模型单实例：仅在model_service进程中加载一次模型，避免GPU和内存浪费
2. 多核心利用：数据生成任务由ProcessPoolExecutor多进程处理，充分发挥Threadripper 2950X的16核心性能
3. 天然同步：队列自动实现推理请求的串行处理，适配单GPU无法并行推理的特性，无需额外锁机制

额外优化建议 #

• 调整队列maxsize参数，避免数据生成速度远快于推理导致内存堆积
• 若TensorFlow版本支持，可在模型服务中使用tf.data.Dataset异步加载请求，进一步提升推理效率
• 数据生成逻辑可拆分为更细粒度的任务，减少单个任务的执行时间，提升结果汇总的实时性

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