Python PyTorch emit_nvtx用法及代碼示例

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

用法:

class torch.autograd.profiler.emit_nvtx(enabled=True, record_shapes=False)

參數：

• enabled(bool,可選的,默認=真) -設置enabled=False 使這個上下文管理器成為no-op。默認值：True。

• record_shapes(bool,可選的,默認=假) -如果 record_shapes=True ，包每個 autograd 操作的 nvtx 範圍將附加有關該操作接收到的張量參數大小的信息，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非張量參數將由 [] 表示。參數將按照後端操作收到的順序列出。請注意，此順序可能與在 Python 端傳遞這些參數的順序不匹配。另請注意，形狀記錄可能會增加創建 nvtx 範圍的開銷。

使每個 autograd 操作發出 NVTX 範圍的上下文管理器。

在 nvprof 下運行程序時很有用：

nvprof --profile-from-start off -o trace_name.prof -- <regular command here>

不幸的是，沒有辦法強製 nvprof 將它收集的數據刷新到磁盤，因此對於 CUDA 分析，必須使用此上下文管理器來注釋 nvprof 跟蹤並等待進程退出，然後再檢查它們。然後，可以使用 NVIDIA Visual Profiler (nvvp) 來可視化時間線，或者 torch.autograd.profiler.load_nvprof() 可以加載結果進行檢查，例如在 Python REPL 中。

示例

>>> with torch.cuda.profiler.profile():
...     model(x) # Warmup CUDA memory allocator and profiler
...     with torch.autograd.profiler.emit_nvtx():
...         model(x)

Forward-backward相關性

在 Nvidia Visual Profiler 中查看使用 emit_nvtx 創建的配置文件時，將每個 backward-pass 操作與相應的 forward-pass 操作相關聯可能很困難。為了簡化此任務，emit_nvtx 將序列號信息附加到它生成的範圍中。

在前向傳遞期間，每個函數範圍都用 seq=<N> 修飾。 seq 是一個運行計數器，每次創建一個新的後向函數對象時遞增，並為後向存儲。因此，與每個前向函數範圍關聯的seq=<N> 注釋告訴您，如果此前向函數創建了一個後向函數對象，則後向對象將接收序列號 N。在後向傳遞期間，包裝每個 C++ 的頂級範圍後向函數的 apply() 調用用 stashed seq=<M> 修飾。 M 是創建後向對象的序列號。通過比較向後的stashed seq 數字和向前的seq 數字，您可以追蹤哪個向前操作創建了每個向後函數。

在向後傳遞期間執行的任何函數也用 seq=<N> 修飾。在默認向後(使用 create_graph=False )期間，此信息無關緊要，事實上，對於所有此類函數，N 可能隻是 0。隻有與後向函數對象的apply() 方法關聯的頂級範圍才有用，作為將這些函數對象與早期前向傳遞相關聯的一種方式。

Double-backward

另一方麵，如果正在進行帶有create_graph=True 的反向傳遞(換句話說，如果您正在設置double-backward)，則反向期間的每個函數的執行都會被賦予一個非零且有用的seq=<N>。這些函數本身可能會創建函數對象，以便稍後在double-backward 期間執行，就像前向傳遞中的原始函數一樣。向後和double-backward 之間的關係在概念上與向前和向後的關係相同：函數仍然發出current-sequence-number-tagged 範圍，它們創建的函數對象仍然存儲這些序列號，並且在最終的double-backward 期間，函數對象的apply() 範圍仍然用stashed seq 數字標記，可以與反向傳遞的seq 數字進行比較。