PyTorch CUDA 内存可视化与快照分析实战用 memory_viz 定位显存泄漏与 OOM【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读本文围绕 PyTorch 官方文档 docs/source/torch_cuda_memory.md 展开系统讲解如何使用torch.cuda.memory._record_memory_history()、_snapshot()与_dump_snapshot()生成 CUDA 内存快照并借助交互式可视化器Active Memory Timeline、Allocator State History逐块还原张量的分配与释放过程。读完本文你将掌握一套可落地的显存排查方法论既能定位“内存被谁占用”也能识别非 PyTorch 分配器管理的显存还能通过快照 API 的底层字段理解分配器 Segment / Block 的真实行为。一、理解 PyTorch CUDA 内存分析的基本原理1.1 快照能记录什么PyTorch 提供了一套内存快照机制在任意时刻记录 CUDA 已分配内存的状态并可选择性地记录导致该状态的分配事件历史。生成的快照可以拖拽到交互式查看器pytorch.org/memory_viz中浏览查看器是运行在本地的 JavaScript 应用不会上传任何快照数据适合在含有敏感模型信息的场景下使用。需要特别明确快照的可见范围来自原文档的 note 及源码实现默认情况下内存分析器只能看到经PyTorch 分配器分配和管理的 CUDA 设备内存例如torch.empty(..., devicecuda)或通过torch.cuda.memory.CUDAPluggableAllocator分配的内存CPUPinned 内存主机锁页内存可以通过向_record_memory_history传入record_pinned_host_memoryTrue选择性纳入直接在 C 中调用 CUDA API如cudaMalloc、cuMemCreate分配的内存或通过第三方 Python 绑定如 cuda-python分配的内存不会出现在 PyTorch 内存分析器中NCCL是典型代表——它用于 CUDA 设备上的分布式通信但其 GPU 内存分配绕过了 PyTorch 分配器因此不会被快照统计。1.2 快照字典的底层结构从源码 torch/cuda/memory.py 可以看到_snapshot()返回的字典具有如下 TypedDict 结构理解它是读懂后续可视化视图的前提segmentsSegment列表每个 Segment 是一次cudaMalloc返回的内存所有 Segment 大小之和即为“保留内存reserved memory”总量Segment 会被缓存复用若复用尺寸小于 Segment则会被拆分为多个 Blocktorch.cuda.empty_cache()只会释放完全空闲的 Segmentdevice_tracesTraceEntry列表仅在开启内存历史记录时存在逐条记录分配器动作Segment关键字段address、total_sizecudaMalloc 分配大小、stream、segment_typesmall或large1MB 为 large、segment_pool_id、allocated_size、active_size、blocksBlock关键字段size、requested_size申请大小可能因对齐小于size、address、stateactive_allocated/active_awaiting_free/inactive、frames分配发生处的栈回溯TraceEntry关键字段actionalloc、free_requested、free_completed、segment_alloc、segment_free、oom、snapshot、annotate、addr、frames、size、stream、pool_idOOM 事件还包含device_freeOOM 发生时 CUDA 仍报告的空闲量。这一结构意味着你不仅能看到“谁占用了内存”还能精确区分“保留但未使用”inactive与“真正被张量持有”active_allocated的内存。二、生成一份内存快照三行代码的标准流程原文档给出了生成快照的经典模式——先开启内存历史再运行待观察代码最后把快照 pickle 保存到文件# 开启内存历史会为快照附加 traceback 与事件历史 torch.cuda.memory._record_memory_history() run_your_code() torch.cuda.memory._dump_snapshot(my_snapshot.pickle)保存后打开pytorch.org/memory_viz把 pickle 文件拖拽进查看器即可开始分析。2.1_record_memory_history的参数全景该函数在源码中的完整签名torch/cuda/memory.py远不止enabled一个参数以下是各参数含义与默认值参数可选值 / 默认值作用enabledNone/state/all默认allNone关闭记录state仅保留当前已分配内存的信息all额外保留全部 alloc/free 调用历史contextNone/state/alloc/all默认allNone不记录任何 tracebackstate记录当前已分配内存的 tracebackalloc额外记录 alloc 调用的 tracebackall再额外记录 free 调用的 tracebackstackspython/all默认allpython仅包含 Python、TorchScript 与 inductor 帧all额外包含 C 帧max_entriesint默认sys.maxsize历史中最多保留的 alloc/free 事件数长时运行任务务必设置合理上限每条记录约几 KB过大将导致内存爆炸详见下文环形缓冲clear_historybool默认False开启时是否清空历史skip_actionslist[str]跳过记录的动作类型用于降低内存开销可选alloc、free_requested、free_completed、segment_alloc、segment_free、oom、snapshotrecord_pinned_host_memorybool默认False同时记录 CPU 锁页内存分配器的历史数据位于快照的host_segments与host_traces键record_cudabool默认True是否记录 CUDA 设备分配器置False可与record_pinned_host_memoryTrue组合实现只记录主机锁页内存2.2 环形缓冲与开销生产环境开启的取舍源码注释揭示了历史记录的环形缓冲实现torch/cuda/memory.pyif (record_history) { if (alloc_trace-size() alloc_trace_max_entries_) { alloc_trace-emplace_back(te); } else { (*alloc_trace)[alloc_trace_next] te; if (alloc_trace_next alloc_trace_max_entries_) { alloc_trace_next 0; } } }当达到max_entries上限后新事件会覆盖最旧事件因此长时间运行的任务只保留最近一段历史延迟开销较低Python 回溯收集约 2µs/条C 回溯收集约 50ns/帧多数程序折算约 2µs/条因此官方建议即使生产任务也可开启以备日后排查快照文件体积_dump_snapshot的 docstring 明确指出文件大小随max_entries与每条栈深增长每条记录数 KB长时运行 大max_entries时文件可达 GB 级。2.3 配套 API_snapshot与_dump_snapshot_snapshot(deviceNone, augment_with_fx_tracesFalse)返回内存快照字典。device为None时捕获当前设备augment_with_fx_tracesTrue时会为栈帧附加 FX 调试信息fx_node_op、fx_node_name、fx_original_trace把生成的 FX 代码映射回原始模型源码_dump_snapshot(filenamedump_snapshot.pickle, augment_with_fx_tracesFalse)把_snapshot()结果 pickle 到文件供pytorch.org/memory_viz打开torch/cuda/memory.py。三、包含 Pinned主机内存锁页内存分析默认快照只含 CUDA 设备内存。若要一并捕获 CPU 锁页内存分配例如pin_memoryTrue创建的张量传入record_pinned_host_memoryTruetorch.cuda.memory._record_memory_history(record_pinned_host_memoryTrue) run_your_code() snapshot torch.cuda.memory._snapshot() # 主机分配器数据位于 snapshot[host_segments] 与 snapshot[host_traces] torch.cuda.memory._dump_snapshot(my_snapshot.pickle)若只想记录主机锁页内存、完全跳过 CUDA 设备分配器可将record_pinned_host_memoryTrue与record_cudaFalse组合torch.cuda.memory._record_memory_history(record_pinned_host_memoryTrue, record_cudaFalse)注意当前pytorch.org/memory_viz可视化器尚不支持展示主机内存数据需要以编程方式从快照字典中检查host_segments与host_traces键对应源码 torch/cuda/memory.py 的参数说明。四、Active Memory Timeline查看张量生命周期Active Memory Timeline活跃内存时间线展示快照内某个 GPU 上所有存活张量随时间的变化是定位“某块显存什么时候被谁持有”的第一站可对绘图区域进行 Pan/Zoom聚焦较小的分配块鼠标悬停已分配块可查看该块被分配时的栈回溯以及地址等细节数据量较大时可调整 detail 滑块减少渲染的分配数量以提升性能。下图即该视图的典型形态图片来自 docs/source/_static/img/torch_cuda_memory/active_memory_timeline.png五、Allocator State History还原每个分配事件的分配器状态Allocator State History分配器状态历史在左侧以时间线列出单个分配器事件。选中某个事件后右侧会给出该事件发生时分配器状态的视觉汇总汇总展示每个从cudaMalloc返回的 Segment以及它如何被拆分成“单个分配 Block”或“空闲空间”悬停 Segment / Block 可查看内存分配时的栈回溯悬停事件如张量被释放可查看事件发生时的栈回溯OOM 事件会被标记出来。当一次分配失败但保留内存reserved memory仍然充足时查看 OOM 时刻的内存状态往往能揭示失败根因——例如碎片化导致无法满足大块连续分配。下图即该视图的典型形态图片来自 docs/source/_static/img/torch_cuda_memory/allocator_state_history.png5.1 地址标识字符串的跨视图追踪栈回溯信息还会报告分配发生的地址。形如b7f064c000000_0的字符串含义是block地址为7f064c000000且这是该地址第_0第 0次被分配。这个唯一字符串可以被在 Active Memory Timeline 中检索查看张量存活区间在 Active State History 中搜索检查张量被分配或释放时的内存状态。这构成了“时间线定位分配 → 状态历史还原现场”的完整排查闭环。六、识别非 PyTorch 分配对比 device_memory_used 与快照如果怀疑 CUDA 内存被分配在 PyTorch 分配器之外典型如 NCCL 通信缓冲、cuda-python 绑定分配可以通过pynvml 包收集裸 CUDA 分配信息再与 PyTorch 报告的分配量对比从而确认差异来源。使用device_memory_used获取 PyTorch 之外的原始显存占用import torch device_idx ... print(torch.cuda.device_memory_used(device_idx))从源码实现看torch/cuda/init.py在 CUDA 平台上该函数通过 pynvml 的nvmlDeviceGetMemoryInfo(handle).used返回全局已用显存字节数即整个进程含非 PyTorch 分配的真实占用在 ROCm/HIP 平台torch.version.hip非空上则改用 amdsmi 的amdsmi_get_gpu_vram_usage换算字节数。因此对比公式为device_memory_used()全局真实占用− 快照中 PyTorch 分配的合计 ≈ 非 PyTorch 分配 / 其他进程占用。若差值显著应重点排查 NCCL、自定义 CUDA 扩展与第三方库的直接 CUDA API 调用。七、完整排查工作流从怀疑到定位综合以上能力一套可复用的显存排查流程如下开启历史torch.cuda.memory._record_memory_history()生产长任务建议显式设置max_entries并视需要设置stackspython降低体积运行并转储运行待观察代码段torch.cuda.memory._dump_snapshot(my_snapshot.pickle)时间线初筛在 Active Memory Timeline 中定位“峰值时刻”与异常长寿命分配块状态历史深挖在 Allocator State History 中检索该块的地址标识字符串如b7f064c000000_0查看其分配 / 释放时刻的分配器全局状态若存在 OOM 事件检查该时刻保留内存与空闲块的分布以判断是否碎片化排除外部分配用torch.cuda.device_memory_used()与快照总量对比识别 NCCL 等非 PyTorch 分配必要时补查主机内存以record_pinned_host_memoryTrue重新记录编程检查host_segments/host_traces程序化分析对于超大规模快照可直接在 Python 中操作_snapshot()返回的字典segments / blocks / device_traces自行统计各栈帧累计分配量而不必依赖可视化器。仓库中 torch/cuda/_memory_viz.py 还提供了基于快照生成 flamegraphsegments、memory、compare等函数的命令行辅助工具可用于批量对比前后两次快照的差异适合在 CI 或脚本化巡检中落地。八、相关 API 速查以下 API 均在torch.cuda.memory模块下torch/cuda/memory.py_record_memory_history(enabledall, contextall, stacksall, max_entriessys.maxsize, deviceNone, clear_historyFalse, skip_actionsNone, record_pinned_host_memoryFalse, record_cudaTrue)开启 / 关闭内存历史与栈回溯记录_snapshot(deviceNone, augment_with_fx_tracesFalse)返回当前 CUDA 内存状态字典segments / device_traces以及可选的 host_segments / host_traces_dump_snapshot(filenamedump_snapshot.pickle, augment_with_fx_tracesFalse)将快照 pickle 落盘供pytorch.org/memory_viz加载torch.cuda.device_memory_used(deviceNone)位于 torch/cuda/init.py返回设备全局已用显存NVML 或 AMDSMI用于对比识别非 PyTorch 分配。小结通过快照 可视化器PyTorch 把 CUDA 分配器内部的 Segment / Block 结构与每次 alloc / free 事件完整暴露给开发者Active Memory Timeline 回答“内存何时被占”Allocator State History 回答“分配失败时内存长什么样”device_memory_used对比则划清了 PyTorch 与外部分配的边界。结合本文对底层快照字典结构与各参数开销的说明你可以依据任务规模单次训练 / 长时服务 / 分布式选择合适的记录粒度把显存排查从“盲猜”升级为“有据可查”。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch CUDA 内存可视化与快照分析实战：用 memory_viz 定位显存泄漏与 OOM
PyTorch CUDA 内存可视化与快照分析实战用 memory_viz 定位显存泄漏与 OOM【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读本文围绕 PyTorch 官方文档 docs/source/torch_cuda_memory.md 展开系统讲解如何使用torch.cuda.memory._record_memory_history()、_snapshot()与_dump_snapshot()生成 CUDA 内存快照并借助交互式可视化器Active Memory Timeline、Allocator State History逐块还原张量的分配与释放过程。读完本文你将掌握一套可落地的显存排查方法论既能定位“内存被谁占用”也能识别非 PyTorch 分配器管理的显存还能通过快照 API 的底层字段理解分配器 Segment / Block 的真实行为。一、理解 PyTorch CUDA 内存分析的基本原理1.1 快照能记录什么PyTorch 提供了一套内存快照机制在任意时刻记录 CUDA 已分配内存的状态并可选择性地记录导致该状态的分配事件历史。生成的快照可以拖拽到交互式查看器pytorch.org/memory_viz中浏览查看器是运行在本地的 JavaScript 应用不会上传任何快照数据适合在含有敏感模型信息的场景下使用。需要特别明确快照的可见范围来自原文档的 note 及源码实现默认情况下内存分析器只能看到经PyTorch 分配器分配和管理的 CUDA 设备内存例如torch.empty(..., devicecuda)或通过torch.cuda.memory.CUDAPluggableAllocator分配的内存CPUPinned 内存主机锁页内存可以通过向_record_memory_history传入record_pinned_host_memoryTrue选择性纳入直接在 C 中调用 CUDA API如cudaMalloc、cuMemCreate分配的内存或通过第三方 Python 绑定如 cuda-python分配的内存不会出现在 PyTorch 内存分析器中NCCL是典型代表——它用于 CUDA 设备上的分布式通信但其 GPU 内存分配绕过了 PyTorch 分配器因此不会被快照统计。1.2 快照字典的底层结构从源码 torch/cuda/memory.py 可以看到_snapshot()返回的字典具有如下 TypedDict 结构理解它是读懂后续可视化视图的前提segmentsSegment列表每个 Segment 是一次cudaMalloc返回的内存所有 Segment 大小之和即为“保留内存reserved memory”总量Segment 会被缓存复用若复用尺寸小于 Segment则会被拆分为多个 Blocktorch.cuda.empty_cache()只会释放完全空闲的 Segmentdevice_tracesTraceEntry列表仅在开启内存历史记录时存在逐条记录分配器动作Segment关键字段address、total_sizecudaMalloc 分配大小、stream、segment_typesmall或large1MB 为 large、segment_pool_id、allocated_size、active_size、blocksBlock关键字段size、requested_size申请大小可能因对齐小于size、address、stateactive_allocated/active_awaiting_free/inactive、frames分配发生处的栈回溯TraceEntry关键字段actionalloc、free_requested、free_completed、segment_alloc、segment_free、oom、snapshot、annotate、addr、frames、size、stream、pool_idOOM 事件还包含device_freeOOM 发生时 CUDA 仍报告的空闲量。这一结构意味着你不仅能看到“谁占用了内存”还能精确区分“保留但未使用”inactive与“真正被张量持有”active_allocated的内存。二、生成一份内存快照三行代码的标准流程原文档给出了生成快照的经典模式——先开启内存历史再运行待观察代码最后把快照 pickle 保存到文件# 开启内存历史会为快照附加 traceback 与事件历史 torch.cuda.memory._record_memory_history() run_your_code() torch.cuda.memory._dump_snapshot(my_snapshot.pickle)保存后打开pytorch.org/memory_viz把 pickle 文件拖拽进查看器即可开始分析。2.1_record_memory_history的参数全景该函数在源码中的完整签名torch/cuda/memory.py远不止enabled一个参数以下是各参数含义与默认值参数可选值 / 默认值作用enabledNone/state/all默认allNone关闭记录state仅保留当前已分配内存的信息all额外保留全部 alloc/free 调用历史contextNone/state/alloc/all默认allNone不记录任何 tracebackstate记录当前已分配内存的 tracebackalloc额外记录 alloc 调用的 tracebackall再额外记录 free 调用的 tracebackstackspython/all默认allpython仅包含 Python、TorchScript 与 inductor 帧all额外包含 C 帧max_entriesint默认sys.maxsize历史中最多保留的 alloc/free 事件数长时运行任务务必设置合理上限每条记录约几 KB过大将导致内存爆炸详见下文环形缓冲clear_historybool默认False开启时是否清空历史skip_actionslist[str]跳过记录的动作类型用于降低内存开销可选alloc、free_requested、free_completed、segment_alloc、segment_free、oom、snapshotrecord_pinned_host_memorybool默认False同时记录 CPU 锁页内存分配器的历史数据位于快照的host_segments与host_traces键record_cudabool默认True是否记录 CUDA 设备分配器置False可与record_pinned_host_memoryTrue组合实现只记录主机锁页内存2.2 环形缓冲与开销生产环境开启的取舍源码注释揭示了历史记录的环形缓冲实现torch/cuda/memory.pyif (record_history) { if (alloc_trace-size() alloc_trace_max_entries_) { alloc_trace-emplace_back(te); } else { (*alloc_trace)[alloc_trace_next] te; if (alloc_trace_next alloc_trace_max_entries_) { alloc_trace_next 0; } } }当达到max_entries上限后新事件会覆盖最旧事件因此长时间运行的任务只保留最近一段历史延迟开销较低Python 回溯收集约 2µs/条C 回溯收集约 50ns/帧多数程序折算约 2µs/条因此官方建议即使生产任务也可开启以备日后排查快照文件体积_dump_snapshot的 docstring 明确指出文件大小随max_entries与每条栈深增长每条记录数 KB长时运行 大max_entries时文件可达 GB 级。2.3 配套 API_snapshot与_dump_snapshot_snapshot(deviceNone, augment_with_fx_tracesFalse)返回内存快照字典。device为None时捕获当前设备augment_with_fx_tracesTrue时会为栈帧附加 FX 调试信息fx_node_op、fx_node_name、fx_original_trace把生成的 FX 代码映射回原始模型源码_dump_snapshot(filenamedump_snapshot.pickle, augment_with_fx_tracesFalse)把_snapshot()结果 pickle 到文件供pytorch.org/memory_viz打开torch/cuda/memory.py。三、包含 Pinned主机内存锁页内存分析默认快照只含 CUDA 设备内存。若要一并捕获 CPU 锁页内存分配例如pin_memoryTrue创建的张量传入record_pinned_host_memoryTruetorch.cuda.memory._record_memory_history(record_pinned_host_memoryTrue) run_your_code() snapshot torch.cuda.memory._snapshot() # 主机分配器数据位于 snapshot[host_segments] 与 snapshot[host_traces] torch.cuda.memory._dump_snapshot(my_snapshot.pickle)若只想记录主机锁页内存、完全跳过 CUDA 设备分配器可将record_pinned_host_memoryTrue与record_cudaFalse组合torch.cuda.memory._record_memory_history(record_pinned_host_memoryTrue, record_cudaFalse)注意当前pytorch.org/memory_viz可视化器尚不支持展示主机内存数据需要以编程方式从快照字典中检查host_segments与host_traces键对应源码 torch/cuda/memory.py 的参数说明。四、Active Memory Timeline查看张量生命周期Active Memory Timeline活跃内存时间线展示快照内某个 GPU 上所有存活张量随时间的变化是定位“某块显存什么时候被谁持有”的第一站可对绘图区域进行 Pan/Zoom聚焦较小的分配块鼠标悬停已分配块可查看该块被分配时的栈回溯以及地址等细节数据量较大时可调整 detail 滑块减少渲染的分配数量以提升性能。下图即该视图的典型形态图片来自 docs/source/_static/img/torch_cuda_memory/active_memory_timeline.png五、Allocator State History还原每个分配事件的分配器状态Allocator State History分配器状态历史在左侧以时间线列出单个分配器事件。选中某个事件后右侧会给出该事件发生时分配器状态的视觉汇总汇总展示每个从cudaMalloc返回的 Segment以及它如何被拆分成“单个分配 Block”或“空闲空间”悬停 Segment / Block 可查看内存分配时的栈回溯悬停事件如张量被释放可查看事件发生时的栈回溯OOM 事件会被标记出来。当一次分配失败但保留内存reserved memory仍然充足时查看 OOM 时刻的内存状态往往能揭示失败根因——例如碎片化导致无法满足大块连续分配。下图即该视图的典型形态图片来自 docs/source/_static/img/torch_cuda_memory/allocator_state_history.png5.1 地址标识字符串的跨视图追踪栈回溯信息还会报告分配发生的地址。形如b7f064c000000_0的字符串含义是block地址为7f064c000000且这是该地址第_0第 0次被分配。这个唯一字符串可以被在 Active Memory Timeline 中检索查看张量存活区间在 Active State History 中搜索检查张量被分配或释放时的内存状态。这构成了“时间线定位分配 → 状态历史还原现场”的完整排查闭环。六、识别非 PyTorch 分配对比 device_memory_used 与快照如果怀疑 CUDA 内存被分配在 PyTorch 分配器之外典型如 NCCL 通信缓冲、cuda-python 绑定分配可以通过pynvml 包收集裸 CUDA 分配信息再与 PyTorch 报告的分配量对比从而确认差异来源。使用device_memory_used获取 PyTorch 之外的原始显存占用import torch device_idx ... print(torch.cuda.device_memory_used(device_idx))从源码实现看torch/cuda/init.py在 CUDA 平台上该函数通过 pynvml 的nvmlDeviceGetMemoryInfo(handle).used返回全局已用显存字节数即整个进程含非 PyTorch 分配的真实占用在 ROCm/HIP 平台torch.version.hip非空上则改用 amdsmi 的amdsmi_get_gpu_vram_usage换算字节数。因此对比公式为device_memory_used()全局真实占用− 快照中 PyTorch 分配的合计 ≈ 非 PyTorch 分配 / 其他进程占用。若差值显著应重点排查 NCCL、自定义 CUDA 扩展与第三方库的直接 CUDA API 调用。七、完整排查工作流从怀疑到定位综合以上能力一套可复用的显存排查流程如下开启历史torch.cuda.memory._record_memory_history()生产长任务建议显式设置max_entries并视需要设置stackspython降低体积运行并转储运行待观察代码段torch.cuda.memory._dump_snapshot(my_snapshot.pickle)时间线初筛在 Active Memory Timeline 中定位“峰值时刻”与异常长寿命分配块状态历史深挖在 Allocator State History 中检索该块的地址标识字符串如b7f064c000000_0查看其分配 / 释放时刻的分配器全局状态若存在 OOM 事件检查该时刻保留内存与空闲块的分布以判断是否碎片化排除外部分配用torch.cuda.device_memory_used()与快照总量对比识别 NCCL 等非 PyTorch 分配必要时补查主机内存以record_pinned_host_memoryTrue重新记录编程检查host_segments/host_traces程序化分析对于超大规模快照可直接在 Python 中操作_snapshot()返回的字典segments / blocks / device_traces自行统计各栈帧累计分配量而不必依赖可视化器。仓库中 torch/cuda/_memory_viz.py 还提供了基于快照生成 flamegraphsegments、memory、compare等函数的命令行辅助工具可用于批量对比前后两次快照的差异适合在 CI 或脚本化巡检中落地。八、相关 API 速查以下 API 均在torch.cuda.memory模块下torch/cuda/memory.py_record_memory_history(enabledall, contextall, stacksall, max_entriessys.maxsize, deviceNone, clear_historyFalse, skip_actionsNone, record_pinned_host_memoryFalse, record_cudaTrue)开启 / 关闭内存历史与栈回溯记录_snapshot(deviceNone, augment_with_fx_tracesFalse)返回当前 CUDA 内存状态字典segments / device_traces以及可选的 host_segments / host_traces_dump_snapshot(filenamedump_snapshot.pickle, augment_with_fx_tracesFalse)将快照 pickle 落盘供pytorch.org/memory_viz加载torch.cuda.device_memory_used(deviceNone)位于 torch/cuda/init.py返回设备全局已用显存NVML 或 AMDSMI用于对比识别非 PyTorch 分配。小结通过快照 可视化器PyTorch 把 CUDA 分配器内部的 Segment / Block 结构与每次 alloc / free 事件完整暴露给开发者Active Memory Timeline 回答“内存何时被占”Allocator State History 回答“分配失败时内存长什么样”device_memory_used对比则划清了 PyTorch 与外部分配的边界。结合本文对底层快照字典结构与各参数开销的说明你可以依据任务规模单次训练 / 长时服务 / 分布式选择合适的记录粒度把显存排查从“盲猜”升级为“有据可查”。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考