讲师中心 微信公众号
AI工具推荐 视频效率加速

Python项目中PyTorch张量跨进程共享时如何避免额外拷贝开销?

小丽吖_5949

小丽吖_5949

发布时间:2026-10-30 10:17:30

|

302人浏览过

|

来源于php中文网

原创

CPU张量调用share_memory_()并配合spawn启动方式可安全共享;CUDA张量无法跨进程共享,必须使用torch.distributed+NCCL替代。

python项目中pytorch张量跨进程共享时如何避免额外拷贝开销?

CPU张量用 share_memory_() 是安全的,但必须配合 spawn 启动方式

PyTorch 的 share_memory_() 只对 CPU 张量生效,调用后张量底层 Storage 会切换到共享内存区(如 /dev/shm),子进程可直接读写同一块物理内存。但前提是主进程必须用 mp.set_start_method('spawn'),不能用 'fork' —— 后者会复制整个地址空间,导致每个子进程拿到的是独立副本,看似共享实则毫无意义。

常见错误是没设启动方式,或者在 Windows/macOS 上误用 'fork'(不支持)。另外,share_memory_() 返回的是原张量对象,不是新张量,所以要确保在 fork/spawn 前调用,否则子进程中调用无效。

  • 必须在创建进程前调用 tensor.share_memory_()
  • 必须显式设置 mp.set_start_method('spawn')(Linux/macOS/Windows 均兼容)
  • 共享后仍需手动同步:多个进程并发写同一位置无原子性保障,需加锁或按索引分区访问

CUDA张量无法跨进程共享,别碰 share_memory_() 或 pin_memory()

对 tensor.cuda() 调用 share_memory_() 不报错,但静默失败 —— 底层 Storage 仍是独占的 GPU 显存,子进程根本访问不到。同理,pin_memory() 只影响 Host 内存页锁定,和 GPU 显存无关,也不能让 CUDA 张量变“可共享”。

本质限制在于:CUDA 显存地址空间是进程私有的,GPU 驱动不提供跨进程共享显存的机制(除非用 NVLINK + MPS 或 CUDA IPC,但 PyTorch 未封装)。所以所有试图“零拷贝共享 GPU 张量”的方案,最终都会退化为 CPU 中转(如通过 Queue 或 Manager 传递时触发 .cpu())。

立即学习“Python免费学习笔记(深入)”;

macOS UI 自动化技能包 —— 用 Python 控制你的 Mac
macOS UI 自动化技能包 —— 用 Python 控制你的 Mac

让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac,无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启:辅助功能 + 屏幕录制权限

下载
  • tensor.cuda().share_memory_() 无效果,tensor.is_shared() 仍返回 False
  • torch.multiprocessing.Queue.put(tensor.cuda()) 必然触发隐式 .cpu() 和 pickle,显存数据先拷到 CPU 再序列化
  • 若必须多进程用 GPU 数据,唯一低开销路径是:每个进程自己加载/生成张量,或用 torch.distributed + NCCL

替代方案:用 torch.distributed 替代 torch.multiprocessing 处理 GPU 张量

当目标是多进程协同计算 GPU 张量(比如分布式训练、多卡推理),torch.multiprocessing 是错的方向。正确路径是走 torch.distributed,它基于 NCCL(多卡)或 Gloo(CPU/单卡)后端,直接在 GPU 显存间通信,不经过主机内存中转。

例如 dist.broadcast() 或 dist.all_reduce(),输入输出都是 CUDA 张量,NCCL 会在驱动层调度 GPU-to-GPU DMA 传输。这比任何用户态共享机制都更底层、更高效。

  • 避免自己实现跨进程 GPU 共享逻辑,PyTorch 官方不支持且无计划支持
  • torch.distributed.init_process_group(backend='nccl') 是多卡训练事实标准,不是“高级选项”而是必需路径
  • 单机多进程 + 单卡场景下,若坚持用 multiprocessing,只能接受每个进程独占一份 GPU 张量副本

容易被忽略的坑:from_numpy() 和 as_tensor() 在多进程里可能意外破坏共享语义

如果你在主进程中用 np.array 构建数据,再用 torch.as_tensor(arr) 转成张量并 share_memory_(),看起来没问题 —— 但一旦某个子进程调用 tensor.numpy() 或 tensor.data_ptr(),就可能触发 NumPy 的 copy-on-write 机制,导致底层内存被悄悄复制,共享失效。

更隐蔽的是,某些数据增强操作(如 torchvision.transforms 中的归一化)内部会调用 .contiguous() 或 .clone(),直接切断与原始 NumPy 数组的内存连接。

  • 跨进程共享前,确认张量没有被任何中间库“偷偷 clone”过
  • 避免在共享张量上直接调用 .numpy()、.data_ptr() 或依赖连续内存的 API
  • 最稳妥做法:共享原始 NumPy 数组(arr = np.memmap 或 shared_memory.SharedMemory),各进程自行 as_tensor(),而非共享 Tensor 对象本身
共享 GPU 显存这件事,在当前 PyTorch 版本(2.5+)和 CUDA 生态下,不存在用户态零拷贝方案。所有绕过 torch.distributed 的尝试,最终都会撞上驱动层隔离墙。真正要省显存,得从模型并行、梯度检查点、FP16/INT4 量化这些方向入手,而不是在进程共享上死磕。

热门AI工具

更多
UP简历
UP简历 Hot

一款AI办公效率工具,主要用于基于AI技术的免费在线简历制作工具,适合需要提升相关任务效率的用户。

切问学术

切问学术是一款AI论文写作工具,复旦大学NLP团队推出的AI学术智能体。

PixTV
PixTV Hot

PixTV是一款面向AIGC内容创作的AI视频生成工具。

豆包大模型

豆包大模型是一款由字节跳动推出的企业级大语言模型服务平台。

DeepSeek

DeepSeek是一款面向对话、写作、编程和推理场景的AI大模型工具。

立刻MV
立刻MV Hot

立刻MV是一款AI文本写作工具,AI 音乐视频(MV)创作工具。

UpDream
UpDream Hot

一款AI视频创作工具,主要用于哔哩哔哩推出的自研AI视频创作工具,适合需要提升相关任务效率的用户。

二狗PPT
二狗PPT Hot

一款AI演示文稿工具,主要用于专为中式职场打造的AI PPT生成工具,适合需要提升相关任务效率的用户。

WorkBuddy

一款AI办公效率工具,主要用于腾讯云推出的AI原生桌面智能体工作台,适合需要提升相关任务效率的用户。

相关专题

更多
python打包成可执行文件
python打包成可执行文件

本专题为大家带来python打包成可执行文件相关的文章,大家可以免费的下载体验。

1831

2023.07.20

python能做什么
python能做什么

python能做的有:可用于开发基于控制台的应用程序、多媒体部分开发、用于开发基于Web的应用程序、使用python处理数据、系统编程等等。本专题为大家提供python相关的各种文章、以及下载和课程。

4904

2023.07.25

format在python中的用法
format在python中的用法

Python中的format是一种字符串格式化方法,用于将变量或值插入到字符串中的占位符位置。通过format方法,我们可以动态地构建字符串,使其包含不同值。php中文网给大家带来了相关的教程以及文章,欢迎大家前来阅读学习。

1789

2023.07.31

python教程
python教程

Python已成为一门网红语言,即使是在非编程开发者当中,也掀起了一股学习的热潮。本专题为大家带来python教程的相关文章,大家可以免费体验学习。

29357

2023.08.03

python环境变量的配置
python环境变量的配置

Python是一种流行的编程语言,被广泛用于软件开发、数据分析和科学计算等领域。在安装Python之后,我们需要配置环境变量,以便在任何位置都能够访问Python的可执行文件。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

3507

2023.08.04

python eval
python eval

eval函数是Python中一个非常强大的函数,它可以将字符串作为Python代码进行执行,实现动态编程的效果。然而,由于其潜在的安全风险和性能问题,需要谨慎使用。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

3507

2023.08.04

scratch和python区别
scratch和python区别

scratch和python的区别:1、scratch是一种专为初学者设计的图形化编程语言,python是一种文本编程语言;2、scratch使用的是基于积木的编程语法,python采用更加传统的文本编程语法等等。本专题为大家提供scratch和python相关的文章、下载、课程内容,供大家免费下载体验。

1223

2023.08.11

python合并两个列表
python合并两个列表

Python是一种强大的编程语言,具有许多方便的功能和工具。在Python中,有多种方法可以合并两个列表。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

616

2023.08.10

Vite 项目如何打包和部署
Vite 项目如何打包和部署

Vite项目打包部署专题汇集生产构建、产物预览、输出目录、资源路径、Nginx与GitHub Pages部署及常见白屏和路由404问题的已审核文章,帮助读者完成上线流程。

40

2026.10.28

热门下载

更多
网站特效
/
网站源码
/
网站素材
/
前端模板

精品课程

更多
相关推荐
/
热门推荐
/
最新课程
Uvicorn官方安装与入门手册
Uvicorn官方安装与入门手册

共0课时 | 0人学习

Python uv官方项目管理指南
Python uv官方项目管理指南

共0课时 | 0人学习

Python uv官方安装与入门文档
Python uv官方安装与入门文档

共0课时 | 0人学习

关于我们 免责申明 举报中心 意见反馈 讲师合作 广告合作 最新更新
php中文网:公益在线php培训,帮助PHP学习者快速成长!
关注服务号
PHP中文网订阅号
每天精选资源文章推送

Copyright 2014-2026 https://www.php.cn/ All Rights Reserved | php.cn