jax.profiler 模組

追蹤和時間效能分析

效能分析計算 描述如何使用 JAX 的追蹤和時間效能分析功能。

start_server(port)	在 port 埠口啟動效能分析器伺服器。
start_trace(log_dir[, create_perfetto_link, ...])	啟動效能分析器追蹤。
stop_trace()	停止目前正在執行的效能分析器追蹤。
trace(log_dir[, create_perfetto_link, ...])	用於進行效能分析器追蹤的情境管理器。
annotate_function(func[, name])	裝飾器，用於為函式的執行產生追蹤事件。
TraceAnnotation	情境管理器，用於在效能分析器中產生追蹤事件。
StepTraceAnnotation(name, **kwargs)	情境管理器，用於在效能分析器中產生步驟追蹤事件。

裝置記憶體效能分析

請參閱 效能分析裝置記憶體 以取得 JAX 裝置記憶體效能分析功能的介紹。

device_memory_profile([backend])	將 JAX 裝置記憶體效能分析擷取為 pprof 格式的協定緩衝區。
save_device_memory_profile(filename[, backend])	收集裝置記憶體效能分析並將其寫入檔案。