jax.profiler.trace#

jax.profiler.trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False)[原始碼]#

用於取得效能分析器追蹤的上下文管理器。

此追蹤將捕捉 CPU、GPU 和/或 TPU 活動，包括 Python 函數和 JAX 裝置端操作。

產生的追蹤可以使用 TensorBoard 檢視。請注意，收集追蹤時不需要執行 TensorBoard。

一次只能收集一個追蹤。如果在另一個追蹤正在執行時開始追蹤，將會引發 RuntimeError。

參數:

  • log_dir (os.PathLike | str) – 儲存效能分析器追蹤的目錄（通常是 TensorBoard 日誌目錄）。

  • create_perfetto_link – 一個布林值，如果為 true，則建立並印出 Perfetto 追蹤檢視器 UI 的連結 (https://ui.perfetto.dev)。程式將會阻擋直到連結開啟且 Perfetto 載入追蹤。

  • create_perfetto_trace – 一個布林值，如果為 true，則額外傾印一個 perfetto_trace.json.gz 檔案，該檔案與 Perfetto 追蹤檢視器 UI 的上傳相容 (https://ui.perfetto.dev)。如果 create_perfetto_link 為 true，也會產生此檔案。如果您想要產生與 Perfetto 相容的追蹤，而無需阻擋程序，這可能會很有用。