如何在multiprocessing.Pool子进程中运行Torch DataLoader并设置num_workers=4

错误原因 #

multiprocessing.Pool创建的进程默认是守护进程（daemon=True），而Python进程规则明确：守护进程不能创建子进程。当PyTorch DataLoader设置num_workers>0时，会启动子进程加速数据读取，这就触发了"daemonic processes are not allowed to have children"错误。

可行解决方案 #

方案一：手动创建非守护进程，控制并发数 #

放弃使用Pool，手动维护进程队列，创建非守护进程执行任务，这样每个任务进程可以正常启动DataLoader的子进程。

import multiprocessing
import torch
from queue import Queue

def process_one_line(line):
    # 加载模型与数据集
    model = load_model()
    dataset = build_dataset(line)
    # 正常设置num_workers=4
    dataloader = torch.utils.data.DataLoader(dataset, num_workers=4)
    
    # 执行推理逻辑示例
    for batch in dataloader:
        output = model(batch)
        # 处理推理结果

def worker_task(queue):
    while True:
        line = queue.get()
        if line is None:  # 终止信号
            break
        process_one_line(line)
        queue.task_done()

if __name__ == '__main__':
    lines = get_lines()  # 10k+ 任务数据
    max_concurrent = 16  # 根据系统资源调整，避免进程过多导致卡顿
    task_queue = Queue(maxsize=max_concurrent)
    
    # 创建非守护进程池
    worker_processes = []
    for _ in range(max_concurrent):
        p = multiprocessing.Process(target=worker_task, args=(task_queue,), daemon=False)
        p.start()
        worker_processes.append(p)
    
    # 提交所有任务
    for line in lines:
        task_queue.put(line)
    
    # 等待所有任务完成
    task_queue.join()
    
    # 发送终止信号并回收进程
    for _ in range(max_concurrent):
        task_queue.put(None)
    for p in worker_processes:
        p.join()

关键注意事项 #

• 控制总进程数：每个任务进程+4个DataLoader worker，总进程数=并发任务数*(1+num_workers)。需根据CPU核心数调整，比如32核CPU建议设置max_concurrent=8，总进程数8*5=40，避免资源过载。
• 模型加载优化：每个子进程重复加载模型会占用大量内存，可考虑：
  • 主进程预加载模型，通过multiprocessing.Manager共享（仅适用于CPU推理）；
  • 使用torch.multiprocessing的spawn上下文，配合模型权重文件路径，让子进程按需加载（减少内存拷贝）。
• 数据集线程安全：确保build_dataset函数是线程/进程安全的，避免多进程同时读写导致的数据异常。

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