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可以是統一的參數元組也可以是每個函數一份參數的列表甚至可傳入生成參數的函數用于按不同規?；鶞蕼y試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 文件命名規則從而在自己的項目中安全、正確地接入這套剖析工具?！久赓M下載鏈接】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可以是統一的參數元組也可以是每個函數一份參數的列表甚至可傳入生成參數的函數用于按不同規?；鶞蕼y試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 文件命名規則從而在自己的項目中安全、正確地接入這套剖析工具?！久赓M下載鏈接】pytorch_geometricGraph Neural Network Library for PyTorch項目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric創作聲明：本文部分內容由AI輔助生成（AIGC），僅供參考