尧图网页设计 
PyTorch Geometric 性能剖析完全指南torch_geometric.profile 模块源码级解析【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric导读训练一个 GNN 模型时运行了多久、显存峰值多少、哪个算子最耗时往往是调优路上最先遇到的问题。PyTorch GeometricPyG在torch_geometric.profile模块中内置了一套完整的性能剖析工具箱从轻量级计时器timeit、可返回结构化显存统计的profileit装饰器、逐层剖析的Profiler到生成 Chrome Trace 的torch_profile、函数级对比基准benchmark与 NVTX 标注工具nvtxit覆盖了 GNN 训练/推理性能分析的全流程。阅读本文后你将掌握这套工具的全部 API、底层实现原理对应源码文件位置以及它们在 PyG 官方 benchmark 脚本中的真实用法能够直接在自己的模型训练脚本中落地使用。本文对应的文档入口为 docs/source/modules/profile.rst该文件通过 Sphinx autosummary 自动展开torch_geometric.profile模块的全部公开成员以下内容即围绕该模块的源码与测试展开。一、模块总览一张 API 地图torch_geometric.profile包的公共接口统一在 torch_geometric/profile/init.py 中声明共 17 个符号分为四类能力类别公开 API实现文件训练/推理计时与显存统计profileit、timeit、get_stats_summary、GPUStats、CUDAStats、GPUStatsSummary、CUDAStatsSummarytorch_geometric/profile/profile.py逐层 Profiler 与 Trace 导出Profiler、torch_profile、xpu_profile、trace_handler、print_time_total、rename_profile_filetorch_geometric/profile/profile.py、torch_geometric/profile/profiler.py模型与数据规模统计count_parameters、get_model_size、get_data_size、get_cpu_memory_from_gc、get_gpu_memory_from_gc、get_gpu_memory_from_nvidia_smi、get_gpu_memory_from_ipextorch_geometric/profile/utils.py函数级对比与 NVTX 标注benchmark、nvtxittorch_geometric/profile/benchmark.py、torch_geometric/profile/nvtx.py其中GPUStats/CUDAStats/GPUStatsSummary/CUDAStatsSummary是profileit与get_stats_summary返回的数据类dataclass也是整个模块的数据骨架会在下文逐一展开。值得注意的是Profiler类位于profiler.py中但未在包的__init__.py导出需要通过from torch_geometric.profile.profiler import Profiler显式导入。二、训练运行时与显存峰值profileit装饰器2.1 基本用法profileit是一个装饰器用于在单次函数调用内同时采集 GPU 运行耗时与显存统计。它要求被装饰函数的第一个参数必须是torch.nn.Module并且只能用于cuda与xpu两种设备源码中会显式抛出AttributeError校验这两点见 profile.py。from torch_geometric.profile import profileit profileit(cuda) def train(model, optimizer, x, edge_index, y): optimizer.zero_grad() out model(x, edge_index) loss criterion(out, y) loss.backward() optimizer.step() return float(loss) loss, stats train(model, optimizer, x, edge_index, y)返回值是一个二元组第一个元素是被装饰函数的原始返回值第二个元素是统计对象——设备为cuda时返回CUDAStats为xpu时返回GPUStats。2.2 底层实现链路从源码看profileit的实现分为三步profile.py推断设备 ID遍历函数入参与关键字参数找到第一个torch.Tensor通过tensor.get_device()推断 GPU 编号若找不到张量或get_device()返回-1CPU 张量则分别抛出AttributeError与RuntimeError。CUDA 内存剖析在 CUDA 设备上会实例化pytorch_memlab的LineProfiler并将model.forward加入被追踪函数列表line_profiler.add_function(args[0].forward)从而获得前向传播逐行/逐算子的显存占用。计时与统计使用torch_gpu.Event(enable_timingTrue)记录开始与结束事件并synchronize()得到秒级耗时随后通过read_from_memlab读取allocated_bytes.all.peak、reserved_bytes.all.peak、active_bytes.all.peak三个峰值指标单位换算为 MB见 profile.py再调用nvidia-smi查询该 GPU 的空闲/已用显存utils.py。2.3 返回的数据结构profileit返回的对象是 dataclassprofile.py字段含义所属类time本次调用 GPU 耗时秒GPUStats/CUDAStatsmax_allocated_gpu峰值已分配显存MBGPUStats/CUDAStatsmax_reserved_gpu峰值已预留reserved显存MBGPUStats/CUDAStatsmax_active_gpu峰值活跃显存MBGPUStats/CUDAStatsnvidia_smi_free_cuda剖析时该 GPU 的空闲显存MB来自nvidia-smi仅CUDAStatsnvidia_smi_used_cuda剖析时该 GPU 的已用显存MB来自nvidia-smi仅CUDAStats2.4 多次运行求汇总get_stats_summary单次统计易受噪声影响实践中通常多次运行后求汇总。get_stats_summary接收一个GPUStats/CUDAStats列表返回对应的GPUStatsSummary/CUDAStatsSummaryprofile.py公共字段time_mean平均耗时、time_std耗时标准差、max_allocated_gpu、max_reserved_gpu、max_active_gpu各取列表最大值CUDAStatsSummary额外包含min_nvidia_smi_free_cuda最小空闲显存与max_nvidia_smi_used_cuda最大已用显存。下面的用法直接取自 PyG 官方测试 test/profile/test_profile.py用profileit(cuda)装饰训练函数跑 5 个 epoch 但只收集后 3 个前 2 个作为 warm-up再求汇总from torch_geometric.profile import get_stats_summary, profileit stats_list [] for epoch in range(5): _, stats train(model, data.x, data.edge_index, data.y) if epoch 2: # Warm-up stats_list.append(stats) stats_summary get_stats_summary(stats_list) print(stats_summary.time_mean, stats_summary.time_std, stats_summary.max_allocated_gpu)三、轻量级计时器timeit上下文管理器当只需要一段代码的运行时长、不关心显存细节时timeit是最轻的选择。它继承contextlib.ContextDecorator既可用作with语句也可用作装饰器profile.pyfrom torch_geometric.profile import timeit torch.no_grad() def test(model, x, edge_index): return model(x, edge_index) with timeit() as t: z test(model, x, edge_index) time t.duration # 秒参数说明logbool默认True为False时不在控制台打印耗时avg_time_divisorint默认0大于 1 时将总耗时除以该值常用于 for 循环内计算平均耗时。两个实现细节值得注意一是进入与退出with块时都会在 CUDA 可用时调用torch.cuda.synchronize()确保测得的是 GPU 算子真正执行完毕的时间而非排队时间二是退出块后t.duration属性才被赋值测试 test_timeit 专门验证了在块内duration不存在、块外才存在这一行为。此外还提供了reset()方法打印当前耗时并重启计时。四、逐层剖析Profiler类timeit给出整体耗时而Profiler可以告诉你每一层模块及每个算子的耗时与显存。它通过递归遍历模型结构_walk_modules会生成(path, is_leaf, module)三元组路径形如(GCN, conv1, lin)对每个叶子模块的forward打上 hook在其内部用torch.profiler.profile记录事件最后汇总成带缩进的分层表格profiler.py。from torch_geometric.profile.profiler import Profiler with Profiler(model, use_cudatorch.cuda.is_available(), profile_memoryTrue) as prof: out model(x, edge_index) # 打印分层剖析表Module / Self CPU total / CPU total / ... / Number of Calls print(prof)构造函数参数profiler.pymodel待剖析的torch.nn.Moduleenabled默认True为False时整体禁用use_cuda默认False是否剖析 CUDA 执行profile_memory默认False是否同时剖析显存paths默认None预定义路径列表用于只剖析指定子模块如[GCN, GCN/conv1]为None时剖析所有叶子模块。剖析输出为分层树状表每层包含 Self CPU total、CPU total、Self CUDA total、CUDA total、Self CPU Mem、CPU Mem、Self CUDA Mem、CUDA Mem、Number of Calls 等列列的具体展示取决于use_cuda与profile_memory开关见 profiler.py。实现中引用了torchprof的分层分组思路且要求 PyTorch 版本不低于 1.8.1见 profiler.py 的版本检查。退出with块后所有被替换的forward会被还原不会污染模型本身。五、Chrome Trace 导出torch_profile与xpu_profiletorch_profile是一个上下文管理器用一行代码接入 PyTorch 官方 profiler并自动把结果导出为 Chrome Trace 格式profile.pyfrom torch_geometric.profile import torch_profile with torch_profile(): model(data.x, data.edge_index)其行为要点自动检测torch.cuda.is_available()决定 activities 是否包含ProfilerActivity.CUDAexport_chrome_traceTrue默认时剖析结束后调用trace_handler先打印按self_cuda_time_total无 CUDA 时按self_cpu_time_total排序的key_averages()表格再把结果导出为当前工作目录下的timeline.jsonexport_chrome_traceFalse时仅打印统计表不导出文件支持可选的csv_data/write_csvprof参数将 Top 5 最耗时算子的 SELF CPU %、SELF CPU、CPU TOTAL %、CPU TOTAL、CUDA 对应列及调用次数写入 CSV见 save_profile_data。配套的两个辅助函数rename_profile_file(*args)把生成的timeline.json重命名为profile-arg1-arg2....json便于多次剖析时保留多个 trace 文件测试 test_torch_profile 中验证了重命名后profile-test_profile.json存在print_time_total(p)按耗时排序打印 profiler 事件表。XPUIntel GPU设备则使用xpu_profile(export_chrome_traceTrue)基于torch.autograd.profiler_legacy.profile(use_xpuTrue)采集打印按self_xpu_time_total排序的表格并可选导出timeline.jsonprofile.py。该分支在测试 test_xpu_profile 中被覆盖。六、模型与数据规模统计utils工具族这类工具解决的是我的模型多大、数据占多少内存的静态规模问题全部实现在 torch_geometric/profile/utils.py 中API返回值说明count_parameters(model)int统计requires_gradTrue的可训练参数量utils.pyget_model_size(model)int将state_dict保存为临时.pt文件后取其磁盘字节数随即删除临时文件utils.pyget_data_size(data)int递归遍历Data/HeteroData的 stores按numel * element_size计算张量理论内存占用通过data_ptr()去重避免重复计数SparseTensor按其 CSR 表示计utils.pyget_cpu_memory_from_gc()int遍历 Python GC 对象累加所有非 CUDA 张量的字节数utils.pyget_gpu_memory_from_gc(device0)int同上但只统计指定设备上的张量utils.pyget_gpu_memory_from_nvidia_smi(device0, digits2)(free, used)MB通过nvidia-smi --query-gpumemory.free/used --formatcsv查询utils.pyget_gpu_memory_from_ipex(device0, digits2)(allocated, reserved, active)MB通过 Intel Extension for PyTorchipex的memory_stats_as_nested_dict读取 XPU 峰值统计utils.py注意get_gpu_memory_from_nvidia_smi的返回值是 MiB其换算系数为 1.0485见medibyte_to_megabyte且源码注释提醒nvidia-smi报告的占用通常高估了本程序实际使用的显存因为包含缓存预留做精确显存分析时优先采用pytorch_memlab/ ipex 的峰值统计。七、函数级对比基准benchmarkbenchmark用于在相同输入上横向对比多个函数的性能输出一张由tabulate渲染的表格benchmark.pyfrom torch_geometric.profile import benchmark benchmark( funcs[add], args(torch.randn(10), torch.randn(10)), num_steps1, num_warmups1, backwardTrue, )参数与行为funcs待对比的函数列表args可以是统一的参数元组也可以是每个函数一份参数的列表甚至可传入生成参数的函数用于按不同规模基准测试num_steps正式计时的步数num_warmups默认 10预热步数预热阶段的耗时不计入结果两者都必须为正整数否则抛ValueErrorbackward默认False为True时同时测量反向传播耗时——实现中会先对输出求和并backward(out_grad)输出为 tuple/list/dict 时自动汇总各张量per_step默认False为True时报告每步平均耗时否则报告总耗时progress_bar默认False使用tqdm显示进度条func_names默认None自定义显示名缺省时从函数__name__推断。无论是否开启backward每次迭代都会在 CUDA 可用时torch.cuda.synchronize()并用time.perf_counter()计时保证结果的确定性。输出表格形如--------------------------------------- | Name | Forward | Backward | Total | |---------------------------------------| | add | 0.0001s | 0.0002s | 0.0003s | ---------------------------------------该行为在测试 test/profile/test_benchmark.py 中通过捕获标准输出得到验证。八、NVTX 标注nvtxitnvtxit用于为函数添加 NVTXNVIDIA Tools Extension范围标记方便在 NVIDIA Nsight 等 GPU 剖析工具中按函数名查看时间线nvtx.pyfrom torch_geometric.profile import nvtxit nvtxit() def forward(self, x, edge_index): return self.propagate(edge_index, xx) nvtxit(custom_name, n_warmups1, n_iters3) def another_func(...): ...参数name可选标记名称缺省为被装饰函数的__name__n_warmups默认 0开始标记前的预热调用次数n_iters可选需要记录的调用次数缺省记录全部。实现上nvtxit通过torch.cuda.nvtx.range_push(f{name}_{iters_so_far})/range_pop()包裹函数体并利用模块级全局变量CUDA_PROFILE_STARTED配合cudaProfilerStart()/cudaProfilerStop()控制 CUDA profiler 的启停nvtx.py。非 CUDA 环境下直接透传执行不做任何标记。其 warm-up、命名、迭代次数等分支行为均在 test/profile/test_nvtx.py 中有系统测试覆盖。九、实战在 PyG 官方 benchmark 中如何组合使用上述工具在 PyG 自带的基准测试脚本中被大量组合使用是最佳实践参考benchmark/citation/train_eval.py 展示了完整流程常规训练 epoch 直接调用train(...)最后一个 run 的最后一个 epoch用with timeit():包裹以获得稳定的单步训练耗时当profilingTrue时再额外执行一次with torch_profile(): train(...)导出timeline.json供 Chrome Tracing 分析。推理路径run_inference同理且在bf16True时结合torch.bfloat16与 AMP autocast 一起剖析。benchmark/citation/gcn.py、benchmark/points/point_cnn.py 等脚本则调用rename_profile_file(gcn, ...)将每次实验的 trace 文件重命名归档避免覆盖。benchmark/kernel/main_performance.py、benchmark/loader/neighbor_loader.py 同样使用了timeit/torch_profile/rename_profile_file组合用于 kernel 级与数据加载环节的性能剖析。十、小结与选型建议面对不同的性能分析诉求torch_geometric.profile提供了从粗到细、从单机到异构设备的完整工具链只需整体耗时 →timeit秒级最轻量需要耗时 显存峰值且能接受多跑几次取汇总 →profileitget_stats_summaryCUDA 依赖pytorch_memlabXPU 依赖 ipex需要定位到具体层/算子的热点 →Profiler逐层剖析表需要可视化时间线 →torch_profile/xpu_profile导出 Chrome Tracetimeline.json/profile-*.json配合rename_profile_file归档需要横向对比多个实现如不同聚合方式→benchmark需要在 Nsight 中按函数定位 →nvtxit需要静态评估模型/数据规模 →count_parameters、get_model_size、get_data_size及 GC /nvidia-smi/ ipex 内存查询族。结合 test/profile 目录下的测试用例可以进一步确认每个 API 的边界行为如设备限制、warm-up 语义、trace 文件命名规则从而在自己的项目中安全、正确地接入这套剖析工具。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch Geometric 性能剖析完全指南：torch_geometric.profile 模块源码级解析
PyTorch Geometric 性能剖析完全指南torch_geometric.profile 模块源码级解析【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric导读训练一个 GNN 模型时运行了多久、显存峰值多少、哪个算子最耗时往往是调优路上最先遇到的问题。PyTorch GeometricPyG在torch_geometric.profile模块中内置了一套完整的性能剖析工具箱从轻量级计时器timeit、可返回结构化显存统计的profileit装饰器、逐层剖析的Profiler到生成 Chrome Trace 的torch_profile、函数级对比基准benchmark与 NVTX 标注工具nvtxit覆盖了 GNN 训练/推理性能分析的全流程。阅读本文后你将掌握这套工具的全部 API、底层实现原理对应源码文件位置以及它们在 PyG 官方 benchmark 脚本中的真实用法能够直接在自己的模型训练脚本中落地使用。本文对应的文档入口为 docs/source/modules/profile.rst该文件通过 Sphinx autosummary 自动展开torch_geometric.profile模块的全部公开成员以下内容即围绕该模块的源码与测试展开。一、模块总览一张 API 地图torch_geometric.profile包的公共接口统一在 torch_geometric/profile/init.py 中声明共 17 个符号分为四类能力类别公开 API实现文件训练/推理计时与显存统计profileit、timeit、get_stats_summary、GPUStats、CUDAStats、GPUStatsSummary、CUDAStatsSummarytorch_geometric/profile/profile.py逐层 Profiler 与 Trace 导出Profiler、torch_profile、xpu_profile、trace_handler、print_time_total、rename_profile_filetorch_geometric/profile/profile.py、torch_geometric/profile/profiler.py模型与数据规模统计count_parameters、get_model_size、get_data_size、get_cpu_memory_from_gc、get_gpu_memory_from_gc、get_gpu_memory_from_nvidia_smi、get_gpu_memory_from_ipextorch_geometric/profile/utils.py函数级对比与 NVTX 标注benchmark、nvtxittorch_geometric/profile/benchmark.py、torch_geometric/profile/nvtx.py其中GPUStats/CUDAStats/GPUStatsSummary/CUDAStatsSummary是profileit与get_stats_summary返回的数据类dataclass也是整个模块的数据骨架会在下文逐一展开。值得注意的是Profiler类位于profiler.py中但未在包的__init__.py导出需要通过from torch_geometric.profile.profiler import Profiler显式导入。二、训练运行时与显存峰值profileit装饰器2.1 基本用法profileit是一个装饰器用于在单次函数调用内同时采集 GPU 运行耗时与显存统计。它要求被装饰函数的第一个参数必须是torch.nn.Module并且只能用于cuda与xpu两种设备源码中会显式抛出AttributeError校验这两点见 profile.py。from torch_geometric.profile import profileit profileit(cuda) def train(model, optimizer, x, edge_index, y): optimizer.zero_grad() out model(x, edge_index) loss criterion(out, y) loss.backward() optimizer.step() return float(loss) loss, stats train(model, optimizer, x, edge_index, y)返回值是一个二元组第一个元素是被装饰函数的原始返回值第二个元素是统计对象——设备为cuda时返回CUDAStats为xpu时返回GPUStats。2.2 底层实现链路从源码看profileit的实现分为三步profile.py推断设备 ID遍历函数入参与关键字参数找到第一个torch.Tensor通过tensor.get_device()推断 GPU 编号若找不到张量或get_device()返回-1CPU 张量则分别抛出AttributeError与RuntimeError。CUDA 内存剖析在 CUDA 设备上会实例化pytorch_memlab的LineProfiler并将model.forward加入被追踪函数列表line_profiler.add_function(args[0].forward)从而获得前向传播逐行/逐算子的显存占用。计时与统计使用torch_gpu.Event(enable_timingTrue)记录开始与结束事件并synchronize()得到秒级耗时随后通过read_from_memlab读取allocated_bytes.all.peak、reserved_bytes.all.peak、active_bytes.all.peak三个峰值指标单位换算为 MB见 profile.py再调用nvidia-smi查询该 GPU 的空闲/已用显存utils.py。2.3 返回的数据结构profileit返回的对象是 dataclassprofile.py字段含义所属类time本次调用 GPU 耗时秒GPUStats/CUDAStatsmax_allocated_gpu峰值已分配显存MBGPUStats/CUDAStatsmax_reserved_gpu峰值已预留reserved显存MBGPUStats/CUDAStatsmax_active_gpu峰值活跃显存MBGPUStats/CUDAStatsnvidia_smi_free_cuda剖析时该 GPU 的空闲显存MB来自nvidia-smi仅CUDAStatsnvidia_smi_used_cuda剖析时该 GPU 的已用显存MB来自nvidia-smi仅CUDAStats2.4 多次运行求汇总get_stats_summary单次统计易受噪声影响实践中通常多次运行后求汇总。get_stats_summary接收一个GPUStats/CUDAStats列表返回对应的GPUStatsSummary/CUDAStatsSummaryprofile.py公共字段time_mean平均耗时、time_std耗时标准差、max_allocated_gpu、max_reserved_gpu、max_active_gpu各取列表最大值CUDAStatsSummary额外包含min_nvidia_smi_free_cuda最小空闲显存与max_nvidia_smi_used_cuda最大已用显存。下面的用法直接取自 PyG 官方测试 test/profile/test_profile.py用profileit(cuda)装饰训练函数跑 5 个 epoch 但只收集后 3 个前 2 个作为 warm-up再求汇总from torch_geometric.profile import get_stats_summary, profileit stats_list [] for epoch in range(5): _, stats train(model, data.x, data.edge_index, data.y) if epoch 2: # Warm-up stats_list.append(stats) stats_summary get_stats_summary(stats_list) print(stats_summary.time_mean, stats_summary.time_std, stats_summary.max_allocated_gpu)三、轻量级计时器timeit上下文管理器当只需要一段代码的运行时长、不关心显存细节时timeit是最轻的选择。它继承contextlib.ContextDecorator既可用作with语句也可用作装饰器profile.pyfrom torch_geometric.profile import timeit torch.no_grad() def test(model, x, edge_index): return model(x, edge_index) with timeit() as t: z test(model, x, edge_index) time t.duration # 秒参数说明logbool默认True为False时不在控制台打印耗时avg_time_divisorint默认0大于 1 时将总耗时除以该值常用于 for 循环内计算平均耗时。两个实现细节值得注意一是进入与退出with块时都会在 CUDA 可用时调用torch.cuda.synchronize()确保测得的是 GPU 算子真正执行完毕的时间而非排队时间二是退出块后t.duration属性才被赋值测试 test_timeit 专门验证了在块内duration不存在、块外才存在这一行为。此外还提供了reset()方法打印当前耗时并重启计时。四、逐层剖析Profiler类timeit给出整体耗时而Profiler可以告诉你每一层模块及每个算子的耗时与显存。它通过递归遍历模型结构_walk_modules会生成(path, is_leaf, module)三元组路径形如(GCN, conv1, lin)对每个叶子模块的forward打上 hook在其内部用torch.profiler.profile记录事件最后汇总成带缩进的分层表格profiler.py。from torch_geometric.profile.profiler import Profiler with Profiler(model, use_cudatorch.cuda.is_available(), profile_memoryTrue) as prof: out model(x, edge_index) # 打印分层剖析表Module / Self CPU total / CPU total / ... / Number of Calls print(prof)构造函数参数profiler.pymodel待剖析的torch.nn.Moduleenabled默认True为False时整体禁用use_cuda默认False是否剖析 CUDA 执行profile_memory默认False是否同时剖析显存paths默认None预定义路径列表用于只剖析指定子模块如[GCN, GCN/conv1]为None时剖析所有叶子模块。剖析输出为分层树状表每层包含 Self CPU total、CPU total、Self CUDA total、CUDA total、Self CPU Mem、CPU Mem、Self CUDA Mem、CUDA Mem、Number of Calls 等列列的具体展示取决于use_cuda与profile_memory开关见 profiler.py。实现中引用了torchprof的分层分组思路且要求 PyTorch 版本不低于 1.8.1见 profiler.py 的版本检查。退出with块后所有被替换的forward会被还原不会污染模型本身。五、Chrome Trace 导出torch_profile与xpu_profiletorch_profile是一个上下文管理器用一行代码接入 PyTorch 官方 profiler并自动把结果导出为 Chrome Trace 格式profile.pyfrom torch_geometric.profile import torch_profile with torch_profile(): model(data.x, data.edge_index)其行为要点自动检测torch.cuda.is_available()决定 activities 是否包含ProfilerActivity.CUDAexport_chrome_traceTrue默认时剖析结束后调用trace_handler先打印按self_cuda_time_total无 CUDA 时按self_cpu_time_total排序的key_averages()表格再把结果导出为当前工作目录下的timeline.jsonexport_chrome_traceFalse时仅打印统计表不导出文件支持可选的csv_data/write_csvprof参数将 Top 5 最耗时算子的 SELF CPU %、SELF CPU、CPU TOTAL %、CPU TOTAL、CUDA 对应列及调用次数写入 CSV见 save_profile_data。配套的两个辅助函数rename_profile_file(*args)把生成的timeline.json重命名为profile-arg1-arg2....json便于多次剖析时保留多个 trace 文件测试 test_torch_profile 中验证了重命名后profile-test_profile.json存在print_time_total(p)按耗时排序打印 profiler 事件表。XPUIntel GPU设备则使用xpu_profile(export_chrome_traceTrue)基于torch.autograd.profiler_legacy.profile(use_xpuTrue)采集打印按self_xpu_time_total排序的表格并可选导出timeline.jsonprofile.py。该分支在测试 test_xpu_profile 中被覆盖。六、模型与数据规模统计utils工具族这类工具解决的是我的模型多大、数据占多少内存的静态规模问题全部实现在 torch_geometric/profile/utils.py 中API返回值说明count_parameters(model)int统计requires_gradTrue的可训练参数量utils.pyget_model_size(model)int将state_dict保存为临时.pt文件后取其磁盘字节数随即删除临时文件utils.pyget_data_size(data)int递归遍历Data/HeteroData的 stores按numel * element_size计算张量理论内存占用通过data_ptr()去重避免重复计数SparseTensor按其 CSR 表示计utils.pyget_cpu_memory_from_gc()int遍历 Python GC 对象累加所有非 CUDA 张量的字节数utils.pyget_gpu_memory_from_gc(device0)int同上但只统计指定设备上的张量utils.pyget_gpu_memory_from_nvidia_smi(device0, digits2)(free, used)MB通过nvidia-smi --query-gpumemory.free/used --formatcsv查询utils.pyget_gpu_memory_from_ipex(device0, digits2)(allocated, reserved, active)MB通过 Intel Extension for PyTorchipex的memory_stats_as_nested_dict读取 XPU 峰值统计utils.py注意get_gpu_memory_from_nvidia_smi的返回值是 MiB其换算系数为 1.0485见medibyte_to_megabyte且源码注释提醒nvidia-smi报告的占用通常高估了本程序实际使用的显存因为包含缓存预留做精确显存分析时优先采用pytorch_memlab/ ipex 的峰值统计。七、函数级对比基准benchmarkbenchmark用于在相同输入上横向对比多个函数的性能输出一张由tabulate渲染的表格benchmark.pyfrom torch_geometric.profile import benchmark benchmark( funcs[add], args(torch.randn(10), torch.randn(10)), num_steps1, num_warmups1, backwardTrue, )参数与行为funcs待对比的函数列表args可以是统一的参数元组也可以是每个函数一份参数的列表甚至可传入生成参数的函数用于按不同规模基准测试num_steps正式计时的步数num_warmups默认 10预热步数预热阶段的耗时不计入结果两者都必须为正整数否则抛ValueErrorbackward默认False为True时同时测量反向传播耗时——实现中会先对输出求和并backward(out_grad)输出为 tuple/list/dict 时自动汇总各张量per_step默认False为True时报告每步平均耗时否则报告总耗时progress_bar默认False使用tqdm显示进度条func_names默认None自定义显示名缺省时从函数__name__推断。无论是否开启backward每次迭代都会在 CUDA 可用时torch.cuda.synchronize()并用time.perf_counter()计时保证结果的确定性。输出表格形如--------------------------------------- | Name | Forward | Backward | Total | |---------------------------------------| | add | 0.0001s | 0.0002s | 0.0003s | ---------------------------------------该行为在测试 test/profile/test_benchmark.py 中通过捕获标准输出得到验证。八、NVTX 标注nvtxitnvtxit用于为函数添加 NVTXNVIDIA Tools Extension范围标记方便在 NVIDIA Nsight 等 GPU 剖析工具中按函数名查看时间线nvtx.pyfrom torch_geometric.profile import nvtxit nvtxit() def forward(self, x, edge_index): return self.propagate(edge_index, xx) nvtxit(custom_name, n_warmups1, n_iters3) def another_func(...): ...参数name可选标记名称缺省为被装饰函数的__name__n_warmups默认 0开始标记前的预热调用次数n_iters可选需要记录的调用次数缺省记录全部。实现上nvtxit通过torch.cuda.nvtx.range_push(f{name}_{iters_so_far})/range_pop()包裹函数体并利用模块级全局变量CUDA_PROFILE_STARTED配合cudaProfilerStart()/cudaProfilerStop()控制 CUDA profiler 的启停nvtx.py。非 CUDA 环境下直接透传执行不做任何标记。其 warm-up、命名、迭代次数等分支行为均在 test/profile/test_nvtx.py 中有系统测试覆盖。九、实战在 PyG 官方 benchmark 中如何组合使用上述工具在 PyG 自带的基准测试脚本中被大量组合使用是最佳实践参考benchmark/citation/train_eval.py 展示了完整流程常规训练 epoch 直接调用train(...)最后一个 run 的最后一个 epoch用with timeit():包裹以获得稳定的单步训练耗时当profilingTrue时再额外执行一次with torch_profile(): train(...)导出timeline.json供 Chrome Tracing 分析。推理路径run_inference同理且在bf16True时结合torch.bfloat16与 AMP autocast 一起剖析。benchmark/citation/gcn.py、benchmark/points/point_cnn.py 等脚本则调用rename_profile_file(gcn, ...)将每次实验的 trace 文件重命名归档避免覆盖。benchmark/kernel/main_performance.py、benchmark/loader/neighbor_loader.py 同样使用了timeit/torch_profile/rename_profile_file组合用于 kernel 级与数据加载环节的性能剖析。十、小结与选型建议面对不同的性能分析诉求torch_geometric.profile提供了从粗到细、从单机到异构设备的完整工具链只需整体耗时 →timeit秒级最轻量需要耗时 显存峰值且能接受多跑几次取汇总 →profileitget_stats_summaryCUDA 依赖pytorch_memlabXPU 依赖 ipex需要定位到具体层/算子的热点 →Profiler逐层剖析表需要可视化时间线 →torch_profile/xpu_profile导出 Chrome Tracetimeline.json/profile-*.json配合rename_profile_file归档需要横向对比多个实现如不同聚合方式→benchmark需要在 Nsight 中按函数定位 →nvtxit需要静态评估模型/数据规模 →count_parameters、get_model_size、get_data_size及 GC /nvidia-smi/ ipex 内存查询族。结合 test/profile 目录下的测试用例可以进一步确认每个 API 的边界行为如设备限制、warm-up 语义、trace 文件命名规则从而在自己的项目中安全、正确地接入这套剖析工具。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考