jax.profiler.start_trace#

jax.profiler.start_trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False)[source]#

開始效能分析器追蹤。

此追蹤將捕捉 CPU、GPU 和/或 TPU 活動，包括 Python 函式和 JAX 裝置端操作。使用 stop_trace() 結束追蹤並將結果儲存到 log_dir。

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

一次只能收集一個追蹤。如果在另一個追蹤正在執行時呼叫 start_trace()，將會引發 RuntimeError。

參數:

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

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

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

返回類型:

None