評價此頁

record_function#

class torch.autograd.profiler.record_function(name, args=None)[原始碼]#

上下文管理器/函式裝飾器,用於在執行 autograd 剖析器時為程式碼塊/函式新增標籤。僅當啟用了 CPU 活動跟蹤時,標籤才會顯示。

當跟蹤程式碼剖析時,此功能非常有用。

引數
  • name (str) – 分配給程式碼塊的標籤。

  • node_id (int) – 節點的 ID,用於分散式剖析。在

  • cases. (非分散式) –

示例

>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
...     y = x**2
...     with torch.autograd.profiler.record_function(
...         "label-z"
...     ):  # label the block
...         z = y**3
...     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
-----------------------------------  ---------------  ---------------  ---------------
pow                                  60.77%           47.470us         3
mul                                  21.73%           25.465us         2
PowBackward0                         12.03%           121.891us        1
torch::autograd::AccumulateGrad      2.70%            6.324us          1
label-z                              2.13%            12.421us         1
torch::autograd::GraphRoot           0.64%            1.503us          1
-----------------------------------  ---------------  ---------------  ---------------
Self CPU time total: 234.344us
CUDA time total: 0.000us