PyTorch 多进程编程权威指南torch.multiprocessing 共享内存、CUDA 张量传递与进程管理【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读本文以 PyTorch 官方文档 docs/source/multiprocessing.md 为主体结合仓库内 torch/multiprocessing 的实际实现与底层 C 共享内存管理器系统讲解torch.multiprocessing的完整能力从共享内存的底层机制file_descriptor与file_system两种策略、CUDA 张量跨进程传递的最佳实践到spawn多进程管理与错误传播。读完本文你将掌握如何让 Tensor 在进程间零拷贝共享、如何规避 CUDA 共享时的内存泄漏与段错误陷阱以及如何用torch.multiprocessing.spawn可靠地管理多个子进程的生命周期。一、torch.multiprocessing 是什么torch.multiprocessing是 Python 标准库multiprocessing的封装它注册了自定义的 reducer利用共享内存让不同进程中的 Tensor 拥有同一份底层数据的共享视图。核心流程是一旦 Tensor/Storage 被移动到共享内存调用 torch.Tensor.share_memory_就可以在不做任何拷贝的情况下把它发送给其他进程。从 torch/multiprocessing/init.py 的模块文档可以看到其设计哲学该模块 API 与原版 100% 兼容——只需把import multiprocessing改为import torch.multiprocessing所有通过队列发送或通过其他机制共享的 Tensor 都会自动移动到共享内存。由于 API 高度相似PyTorch 并未重复文档化大部分内容而是推荐直接参考 Python 标准库的multiprocessing文档。这意味着你已有的多进程代码几乎可以零改动迁移到 PyTorch 版本从而获得 Tensor 的共享内存传递能力。1.1 内部初始化链路在 torch/multiprocessing/init.py 中from multiprocessing import * # 继承标准库全部 API torch._C._multiprocessing_init() # 初始化 C 侧共享内存支持随后调用init_reductions()定义于 torch/multiprocessing/reductions.py向ForkingPickler注册 Tensor/Storage 的自定义序列化函数。这些 reducer 维护一个SharedCachereductions.py它是一个multiprocessing 句柄 → Storage 弱引用的字典带锁保护并在 fork 后重置锁避免死锁。二、策略管理Strategy management共享策略决定 CPU Tensor 底层数据如何在不同进程间传递。torch.multiprocessing提供三个顶层管理函数定义于 torch/multiprocessing/init.py函数作用get_all_sharing_strategies()返回当前系统支持的全部共享策略集合get_sharing_strategy()返回当前生效的共享策略名称set_sharing_strategy(new_strategy)设置共享策略传入不在支持集合中的值时抛出AssertionError2.1 平台默认值策略的默认选择逻辑写在 torch/multiprocessing/init.pyif sys.platform darwin or sys.platform win32: _sharing_strategy file_system _all_sharing_strategies {file_system} else: _sharing_strategy file_descriptor _all_sharing_strategies {file_descriptor, file_system}Linux 等平台默认file_descriptor支持{file_descriptor, file_system}macOS / Windows只支持file_systemfile_descriptor依赖 UNIX 域套接字传递文件描述符在这些平台不可用。仓库测试 test/distributed/rpc/test_share_memory.py 演示了如何在测试中临时切换到file_system并在结束后恢复原策略可作为日常使用的参考写法。三、共享 CUDA 张量约束与三条最佳实践CUDA 张量跨进程共享仅支持 Python 3且必须使用spawn或forkserver启动方式fork无法安全地继承 CUDA 上下文。这与 CPU 张量有本质区别发送进程必须一直持有原始张量直到接收进程不再需要它为止。引用计数在底层已经实现但需要用户遵守以下最佳实践。3.1 实践一在消费端尽快释放内存消费进程拿到张量后应立即使用并删除否则发送进程必须持续为其保留内存# 好及时释放 x queue.get() # 用 x 做点什么 del x# 差持有到进程结束生产者被迫一直保留 x 的内存 x queue.get() # 用 x 做点什么 # 之后还做其他事生产者不得不一直把 x 留在内存里3.2 实践二让生产者进程运行到所有消费者退出这可以防止生产者提前释放消费者仍在使用中的内存# 生产者 # 发送张量、做其他事 event.wait() # 等待消费者发出完成信号# 消费者 # 接收张量并使用 event.set() # 通知生产者可以结束3.3 实践三不要转发接收到的张量CUDA 张量不能在进程间接力传递否则无法正常工作甚至触发段错误# 这样不行 x queue.get() queue_2.put(x)# 正确做法先创建进程本地副本再转发 x queue.get() x_clone x.clone() queue_2.put(x_clone)# 在同一进程内对同一个队列先 put 再 get极可能导致段错误 queue.put(tensor) x queue.get()3.4 崩溃时的内存风险原文档同时给出两条重要警告若主进程突然退出例如收到信号Python 的multiprocessing有时无法清理子进程中断解释器后可能出现资源泄漏若消费进程因致命信号异常死亡只要发送进程还在运行共享张量就可能被永久保留在内存中。CUDA 事件Event也可以通过 IPC 句柄在进程间共享相关实现见 torch/multiprocessing/reductions.py 的_reduce_event/_rebuild_event。四、共享策略详解file_descriptor 与 file_system本节仅适用于 CPU Tensor——CUDA 张量始终使用 CUDA 的 IPC API这是它们唯一可行的共享方式。4.1 文件描述符策略file_descriptor这是默认策略macOS 与 Windows 除外。工作方式存储一旦移入共享内存就从shm_open获得一个文件描述符并缓存到对象上发送时通过 UNIX 套接字等机制把文件描述符传递给接收进程接收进程同样缓存该描述符并执行mmap从而获得存储数据的共享视图。该策略的代价是共享大量张量时系统会长期保持大量打开的文件描述符。若你的系统对打开文件数有较低上限且无法提升应改用file_system策略。4.2 文件系统策略file_system该策略使用shm_open生成的文件名来标识共享内存区域。好处是不需要在实现中缓存文件描述符坏处是容易产生共享内存泄漏——文件创建后不能立即删除其他进程还需要通过它打开自己的视图一旦进程崩溃或被杀死而没有调用存储析构函数这些文件就会残留在系统中持续占用内存直到系统重启或手动释放。为缓解这个问题torch.multiprocessing会派生一个名为torch_shm_manager的守护进程它把自己隔离出当前进程组跟踪进程组内所有的共享内存分配当所有相关进程退出后它会等待片刻确保不会有新连接然后遍历该组分配的全部共享内存文件发现仍存在的便予以释放。这一机制在 C 侧由 torch/lib/libshm/core.cpp 的start_manager()实现——通过forkpipe派生守护进程并交换握手信息socket handle随后以execl启动名为torch_shm_manager的可执行文件。原文档的结论是该清理方法经过测试对各种故障足够健壮但只要系统限制允许且file_descriptor可用官方仍不推荐切换到file_system。五、派生子进程torch.multiprocessing.spawn手动创建多个Process实例并逐个join的做法在单进程场景下没有问题但面对多个进程时存在隐患串行join意味着进程按顺序终止若第一个进程不退出后续进程的终止将无人察觉同时标准库也没有原生的错误传播机制。torch.multiprocessing.spawn正是为解决这些问题而生它提供错误传播、乱序终止检测并在某个进程出错时主动终止其余进程。5.1 spawn 函数签名spawn定义于 torch/multiprocessing/spawn.pyspawn(fn, args(), nprocs1, joinTrue, daemonFalse, start_methodspawn)参数说明fn子进程入口函数必须定义在模块顶层multiprocessing 要求可 pickle以fn(i, *args)调用i为进程索引args传给fn的参数元组nprocs派生的进程数join是否阻塞等待所有进程结束为True返回None为False返回ProcessContextdaemon子进程的 daemon 标志start_method已废弃参数——该方法总是使用spawn如需其他启动方式请用start_processes()5.2 错误传播与清理机制从 torch/multiprocessing/spawn.py 可以看到每个子进程由_wrap包装执行发生异常时traceback 会被 pickle 写入一个独立的临时错误文件pytorch-errorfile-*.pickle使用文件而非队列是为了避免死锁随后以退出码 1 退出。父进程侧ProcessContext.joinspawn.py通过multiprocessing.connection.wait监听所有进程的 sentinel一旦发现某个进程退出码非零就执行分级清理——先给出grace_period让其他进程自行退出然后发送SIGTERM若再等待一个宽限期仍未退出则升级为SIGKILL。随后根据错误文件是否存在抛出两种异常ProcessRaisedException子进程抛出了异常父进程会收到包含原始 traceback 的完整报错ProcessExitedException子进程因信号或非零退出码终止异常携带exit_code与signal_name信息。两者均继承自ProcessExceptionspawn.py记录error_index与error_pid以便定位出错的进程。5.3 ProcessContext 与 SpawnContext调用spawn(..., joinFalse)会返回ProcessContext其join(timeoutNone, grace_periodNone)方法可多次调用返回True表示所有进程已全部成功结束返回False表示还有进程未结束任一进程失败时抛异常并清理其余进程。SpawnContext是ProcessContext的历史名称自 1.4 版本起已重命名仍保留并发出弃用警告spawn.py。5.4 start_processes更通用的底层 API若需要fork、forkserver等其他启动方式应使用start_processes(fn, args, nprocs, join, daemon, start_method)spawn.py。它在forkserver且设置环境变量TORCH_MP_PARALLEL_START1时通过ThreadPoolExecutor并行启动进程以提升性能见 ENV_VAR_PARALLEL_START。该 API 当前未写入官方文档但在 IPython Notebook 等环境中fork往往比spawn表现更好值得了解。六、配套模块pool、queue 与 reductions原文档为这三个模块预留了文档位见 docs/source/multiprocessing.md它们在仓库中均有完整实现6.1 torch.multiprocessing.pooltorch/multiprocessing/pool.py 提供了Pool类继承自multiprocessing.pool.Pool但改用自己的SimpleQueue从而让 Tensor 以共享内存形式在池的进程间传递而非序列化底层数据。其clean_worker在标准 worker 结束后显式调用gc.collect()确保所有析构函数包括共享内存释放被真正执行。6.2 torch.multiprocessing.queuetorch/multiprocessing/queue.py 中的Queue与SimpleQueue通过ConnectionWrapper在发送时使用ForkingPickler序列化对象——这正是 Tensor 走自定义共享内存 reducer 的关键路径也是set_sharing_strategy等设置能生效的底层原因。6.3 torch.multiprocessing.reductionstorch/multiprocessing/reductions.py 是共享机制的核心包含StorageWeakRef对 Storage 的弱引用配合SharedCache跟踪句柄与 Storage 的对应关系SharedCache带锁与 fork 重置机制的句柄缓存长度超限时自动清理已失效引用rebuild_tensor/_reduce_event等重建函数负责在接收端根据句柄与元数据offset、size、stride、requires_grad恢复张量。七、实战要点总结迁移成本极低把import multiprocessing换成import torch.multiprocessing即可让队列中的 Tensor 自动共享内存无需拷贝。CPU 张量Linux 默认file_descriptor策略高效但占用 fd 较多受限于系统 fd 上限时切换set_sharing_strategy(file_system)并依赖torch_shm_manager守护进程清理泄漏文件。CUDA 张量务必使用spawn/forkserver在消费端及时del、让生产者运行到消费者全部退出、绝不转发收到的张量避免内存滞留与段错误。多进程管理优先使用torch.multiprocessing.spawn而非手写循环join它自带错误传播ProcessRaisedException/ProcessExitedException、乱序终止检测与SIGTERM→SIGKILL分级清理显著提升健壮性。守护进程说明torch_shm_manager由 PyTorch 在需要时自动派生实现见 torch/lib/libshm/core.cpp用户无需手动启动只需了解其内存清理职责即可。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考