Python PyTorch profile用法及代碼示例

本文簡要介紹python語言中 torch.autograd.profiler.profile 的用法。

用法:

class torch.autograd.profiler.profile(enabled=True, *, use_cuda=False, record_shapes=False, with_flops=False, profile_memory=False, with_stack=False, with_modules=False, use_kineto=False, use_cpu=True)

參數：

• enabled(bool,可選的) -將此設置為 False 會使此上下文管理器成為 no-op。

• use_cuda(bool,可選的) -啟用 CUDA 事件的計時以及使用 cudaEvent API。為每個張量操作增加大約 4us 的開銷。

• record_shapes(bool,可選的) -如果設置了形狀記錄，則將收集有關輸入尺寸的信息。這允許人們查看在引擎蓋下使用了哪些維度，並使用 prof.key_averages(group_by_input_shape=True) 對它們進行進一步分組。請注意，形狀記錄可能會扭曲您的分析數據。建議使用帶和不帶形狀記錄的單獨運行來驗證時間。對於最底部的事件(在嵌套函數調用的情況下)，偏斜很可能可以忽略不計。但是對於更高級別的函數，由於形狀收集，總的自身 cpu 時間可能會人為地增加。

• with_flops(bool,可選的) -如果設置了with_flops，分析器將使用運算符的輸入形狀估計 FLOPs(浮點運算)值。這允許人們估計硬件性能。目前，此選項僅適用於矩陣乘法和 2D 卷積運算符。

• profile_memory(bool,可選的) -跟蹤張量內存分配/釋放。

• with_stack(bool,可選的) -記錄操作的源信息(文件和行號)。

• with_modules(bool) -記錄操作的調用棧對應的模塊層次結構(包括函數名)。例如如果模塊 A 的轉發調用的模塊 B 的轉發包含 aten::add 操作，則 aten::add 的模塊層次結構為 A.B 請注意，目前僅對 TorchScript 模型而不是即刻模式模型存在此支持。

• use_kineto(bool,可選的) -實驗性的，啟用使用 Kineto 分析器進行分析。

• use_cpu(bool,可選的) -分析 CPU 事件；設置為 False 需要 use_kineto=True 並且可用於降低 GPU-only 分析的開銷。

管理 autograd 分析器狀態並保存結果摘要的上下文管理器。在底層，它隻是記錄在 C++ 中執行的函數的事件，並將這些事件公開給 Python。您可以將任何代碼包裝到其中，它隻會報告 PyTorch 函數的運行時。注意：探查器是線程本地的，並自動傳播到異步任務中

示例

>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
>>>     for _ in range(100):  # any normal python code, really!
>>>         y = x ** 2
>>          y.backward()
>>> # NOTE: some columns were removed for brevity
>>> print(prof.key_averages().table(sort_by="self_cpu_time_total"))
-----------------------------------  ---------------  ---------------  ---------------
Name                                 Self CPU total   CPU time avg     Number of Calls
-----------------------------------  ---------------  ---------------  ---------------
mul                                  32.048ms         32.048ms         200
pow                                  27.041ms         27.041ms         200
PowBackward0                         9.727ms          55.483ms         100
torch::autograd::AccumulateGrad      9.148ms          9.148ms          100
torch::autograd::GraphRoot           691.816us        691.816us        100
-----------------------------------  ---------------  ---------------  ---------------