单GPU下PyTorch并行运行函数问题及最优方案咨询

1. Python multiprocessing库是否仅支持CPU？ #

不是。multiprocessing可以调度GPU任务，但直接用默认的Pool（基于fork启动方式）和PyTorch GPU任务配合极易出问题——fork操作会复制父进程的CUDA上下文，子进程复用该上下文时会引发死锁、资源冲突，这就是你遇到内核无响应的直接原因。

2. 单GPU下运行PyTorch并行函数的最优方案 #

根据你的场景，推荐以下几种方案，按优先级排序：

方案一：用PyTorch优化的多进程库+spawn启动方式 #

PyTorch自带的torch.multiprocessing专门针对GPU多进程做了兼容优化，配合spawn启动方式（从头初始化子进程，不继承父进程CUDA上下文），能避免大部分冲突：

import torch.multiprocessing as mp

def pytorch_classification(param):
    # 必须在子进程内部导入PyTorch、加载模型/数据
    import torch
    # 执行分类任务逻辑
    model = ...  # 加载你的CNN模型
    data = ...   # 加载对应参数的数据集
    # 推理/训练
    ...

if __name__ == "__main__":
    mp.set_start_method('spawn')
    # 根据显存容量设置进程数（比如你单任务占1/3显存，最多开3个）
    with mp.Pool(processes=2) as pool:
        pool.map(pytorch_classification, [param1, param2])

方案二：标准库multiprocessing改用spawn启动 #

如果坚持用标准库，同样需要强制设置spawn启动方式，代码逻辑和上面一致：

from multiprocessing import Pool, set_start_method

def pytorch_classification(param):
    import torch
    # 任务逻辑
    ...

if __name__ == "__main__":
    set_start_method('spawn')
    with Pool(processes=2) as pool:
        pool.map(pytorch_classification, [param1, param2])

方案三：批量并行（同模型多任务场景） #

如果你的多个分类任务基于同一模型，直接把所有任务的输入打包成大batch，用单进程做批量推理/训练，这是内存占用最低、效率最高的方式——只需加载一次库和模型：

import torch

# 加载一次模型
model = ...
model.cuda()

# 打包所有任务的输入为一个batch
all_data = torch.cat([load_data(param1), load_data(param2)]).cuda()

# 批量执行分类
results = model(all_data)

# 拆分结果对应到各个任务
result1, result2 = results.split([len(data1), len(data2)])

额外优化建议 #

• 控制进程数量：根据单任务显存占用设置，避免显存溢出；
• 子进程及时释放资源：任务完成后执行del model、torch.cuda.empty_cache()，减少RAM和显存占用；
• 避免全局作用域加载模型：所有PyTorch相关的初始化必须放在子进程的任务函数内部（spawn启动的子进程不会继承父进程的全局变量）。

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