計算效能分析#

使用 Perfetto 檢視程式追蹤#

我們可以使用 JAX profiler 來產生 JAX 程式的追蹤，這些追蹤可以使用 Perfetto visualizer 進行視覺化。目前，此方法會阻擋程式執行，直到連結被點擊且 Perfetto UI 載入追蹤。如果您希望在沒有任何互動的情況下取得效能分析資訊，請查看下方的 Tensorboard profiler。

with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
  # Run the operations to be profiled
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

此計算完成後，程式會提示您開啟 ui.perfetto.dev 的連結。當您開啟連結時，Perfetto UI 將載入追蹤檔案並開啟視覺化工具。

Perfetto trace viewer

載入連結後，程式執行將會繼續。連結在開啟一次後即失效，但會重新導向至新的 URL，該 URL 仍然有效。然後您可以點擊 Perfetto UI 中的「Share」按鈕來建立追蹤的永久連結，以便與他人分享。

遠端效能分析#

當分析遠端執行的程式碼（例如在託管 VM 上）時，您需要在連接埠 9001 上建立 SSH 通道，連結才能運作。您可以使用以下命令執行此操作

$ ssh -L 9001:127.0.0.1:9001 <user>@<host>

或者如果您使用 Google Cloud

$ gcloud compute ssh <machine-name> -- -L 9001:127.0.0.1:9001

手動擷取#

除了使用 jax.profiler.trace 以程式方式擷取追蹤之外，您還可以透過呼叫 jax.profiler.start_server(<port>) 在感興趣的腳本中啟動效能分析伺服器。如果您只需要效能分析伺服器在腳本的某一部分保持活動狀態，您可以透過呼叫 jax.profiler.stop_server() 來關閉它。

在腳本執行且效能分析伺服器啟動後，我們可以手動擷取和追蹤，方法是執行

$ python -m jax.collect_profile <port> <duration_in_ms>

預設情況下，產生的追蹤資訊會傾印到暫存目錄中，但可以透過傳入 --log_dir=<選擇的目錄> 來覆寫。此外，預設情況下，程式會提示您開啟 ui.perfetto.dev 的連結。當您開啟連結時，Perfetto UI 將載入追蹤檔案並開啟視覺化工具。透過在命令中傳入 --no_perfetto_link 可以停用此功能。或者，您也可以將 Tensorboard 指向 log_dir 來分析追蹤（請參閱下方的「Tensorboard 效能分析」章節）。

TensorBoard 效能分析#

TensorBoard 的 profiler 可以用來分析 JAX 程式。Tensorboard 是獲取和視覺化程式效能追蹤和設定檔的好方法，包括 GPU 和 TPU 上的活動。最終結果看起來像這樣

TensorBoard profiler example

安裝#

TensorBoard profiler 僅適用於與 TensorFlow 捆綁的 TensorBoard 版本。

pip install tensorflow tensorboard-plugin-profile

如果您已經安裝了 TensorFlow，則只需要安裝 tensorboard-plugin-profile pip 套件。請注意僅安裝一個版本的 TensorFlow 或 TensorBoard，否則您可能會遇到 下方 描述的「重複外掛程式」錯誤。請參閱 https://tensorflow.dev.org.tw/guide/profiler 以取得有關安裝 TensorBoard 的更多資訊。

TensorBoard profiler 的 nightly 版本需要 nightly tensorflow 和 tensorboard

pip install tf-nightly tb-nightly tbp-nightly

程式化擷取#

您可以透過 jax.profiler.start_trace() 和 jax.profiler.stop_trace() 方法，對您的程式碼進行檢測以擷取 profiler 追蹤。呼叫 start_trace() 並指定要將追蹤檔案寫入的目錄。這應該是啟動 TensorBoard 時使用的相同 --logdir 目錄。然後，您可以使用 TensorBoard 來檢視追蹤。

例如，要進行 profiler 追蹤

import jax

jax.profiler.start_trace("/tmp/tensorboard")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

請注意 block_until_ready() 呼叫。我們使用它來確保裝置上的執行被追蹤擷取。請參閱 非同步分派 以了解為什麼這是必要的。

您也可以使用 jax.profiler.trace() 內容管理器，作為 start_trace 和 stop_trace 的替代方案

import jax

with jax.profiler.trace("/tmp/tensorboard"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

要檢視追蹤，請先啟動 TensorBoard（如果尚未啟動）

$ tensorboard --logdir=/tmp/tensorboard
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.5.0 at https://127.0.0.1:6006/ (Press CTRL+C to quit)

在此範例中，您應該能夠在 https://127.0.0.1:6006/ 載入 TensorBoard。您可以使用 --port 旗標指定不同的連接埠。如果要在遠端伺服器上執行 JAX，請參閱下方的 在遠端機器上進行效能分析。

然後，在右上角的下拉式選單中選取「Profile」，或直接前往 https://127.0.0.1:6006/#profile。可用的追蹤會出現在左側的「Runs」下拉式選單中。選取您感興趣的執行，然後在「Tools」下方，選取 trace_viewer。您現在應該會看到執行的時間軸。您可以使用 WASD 鍵來瀏覽追蹤，並點擊或拖曳以選取事件，以便在底部查看更多詳細資訊。有關使用追蹤檢視器的更多詳細資訊，請參閱 這些 TensorFlow 文件。

您也可以使用 memory_viewer、op_profile 和 graph_viewer 工具。

透過 TensorBoard 手動擷取#

以下是從正在執行的程式手動觸發 N 秒追蹤的說明。

  1. 啟動 TensorBoard 伺服器

    tensorboard --logdir /tmp/tensorboard/

    您應該能夠在 https://127.0.0.1:6006/ 載入 TensorBoard。您可以使用 --port 旗標指定不同的連接埠。如果要在遠端伺服器上執行 JAX，請參閱下方的 在遠端機器上進行效能分析。

  2. 在您想要分析的 Python 程式或程序中，在開頭附近新增以下程式碼

    import jax.profiler
jax.profiler.start_server(9999)

    這會啟動 TensorBoard 連接的效能分析伺服器。效能分析伺服器必須在您繼續下一步之前執行。當您完成使用伺服器後，您可以呼叫 jax.profiler.stop_server() 來關閉它。

    如果您想要分析長時間執行的程式片段（例如，長時間的訓練迴圈），您可以將其放在程式的開頭，並照常啟動程式。如果您想要分析簡短的程式（例如，微基準測試），一種選擇是在 IPython shell 中啟動效能分析伺服器，並在下一步中啟動擷取後，使用 %run 執行簡短的程式。另一種選擇是在程式的開頭啟動效能分析伺服器，並使用 time.sleep() 給您足夠的時間來啟動擷取。

  3. 開啟 https://127.0.0.1:6006/#profile，然後點擊左上角的「CAPTURE PROFILE」按鈕。輸入「localhost:9999」作為設定檔服務 URL（這是您在上一步中啟動的效能分析伺服器的位址）。輸入您要分析的毫秒數，然後點擊「CAPTURE」。

  4. 如果您想要分析的程式碼尚未執行（例如，如果您在 Python shell 中啟動了效能分析伺服器），請在擷取執行時執行它。

  5. 擷取完成後，TensorBoard 應該會自動重新整理。（並非所有 TensorBoard 效能分析功能都與 JAX 連接，因此最初看起來可能好像沒有擷取到任何內容。）在左側的「Tools」下方，選取 trace_viewer。

    您現在應該會看到執行的時間軸。您可以使用 WASD 鍵來瀏覽追蹤，並點擊或拖曳以選取事件，以便在底部查看更多詳細資訊。有關使用追蹤檢視器的更多詳細資訊，請參閱 這些 TensorFlow 文件。

    您也可以使用 memory_viewer、op_profile 和 graph_viewer 工具。

加入自訂追蹤事件#

預設情況下，追蹤檢視器中的事件大多是低階內部 JAX 函數。您可以使用 jax.profiler.TraceAnnotation 和 jax.profiler.annotate_function() 在您的程式碼中加入您自己的事件和函數。

疑難排解#

GPU 效能分析#

在 GPU 上執行的程式應該會在追蹤檢視器的頂部附近產生 GPU 串流的追蹤。如果您只看到主機追蹤，請檢查您的程式日誌和/或輸出是否有以下錯誤訊息。

如果您收到類似以下的錯誤訊息：Could not load dynamic library 'libcupti.so.10.1'
完整錯誤訊息

W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcupti.so.10.1'; dlerror: libcupti.so.10.1: cannot open shared object file: No such file or directory
2020-06-12 13:19:59.822799: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1422] function cupti_interface_->Subscribe( &subscriber_, (CUpti_CallbackFunc)ApiCallback, this)failed with error CUPTI could not be loaded or symbol could not be found.

將 libcupti.so 的路徑新增至環境變數 LD_LIBRARY_PATH。（嘗試使用 locate libcupti.so 尋找路徑。）例如

export LD_LIBRARY_PATH=/usr/local/cuda-10.1/extras/CUPTI/lib64/:$LD_LIBRARY_PATH

如果您在執行此操作後仍然收到 Could not load dynamic library 訊息，請檢查 GPU 追蹤是否仍然顯示在追蹤檢視器中。即使一切正常運作，有時也會發生此訊息，因為它會在多個位置尋找 libcupti 程式庫。

如果您收到類似以下的錯誤訊息：failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
完整錯誤訊息

E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1445] function cupti_interface_->EnableCallback( 0 , subscriber_, CUPTI_CB_DOMAIN_DRIVER_API, cbid)failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
2020-06-12 14:31:54.097791: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1487] function cupti_interface_->ActivityDisable(activity)failed with error CUPTI_ERROR_NOT_INITIALIZED

執行以下命令（請注意，這需要重新啟動）

echo 'options nvidia "NVreg_RestrictProfilingToAdminUsers=0"' | sudo tee -a /etc/modprobe.d/nvidia-kernel-common.conf
sudo update-initramfs -u
sudo reboot now

有關更多資訊，請參閱 NVIDIA 關於此錯誤的說明文件。

在遠端機器上進行效能分析#

如果您要分析的 JAX 程式在遠端機器上執行，一種選擇是在遠端機器上執行上述所有指示（特別是在遠端機器上啟動 TensorBoard 伺服器），然後使用 SSH 本機連接埠轉發從本機機器存取 TensorBoard Web UI。使用以下 SSH 命令將預設 TensorBoard 連接埠 6006 從本機轉發到遠端機器

ssh -L 6006:localhost:6006 <remote server address>

或者如果您使用 Google Cloud

$ gcloud compute ssh <machine-name> -- -L 6006:localhost:6006

多個 TensorBoard 安裝#

如果啟動 TensorBoard 失敗並出現類似以下的錯誤訊息：ValueError: Duplicate plugins for name projector

通常是因為安裝了兩個版本的 TensorBoard 和/或 TensorFlow（例如，tensorflow、tf-nightly、tensorboard 和 tb-nightly pip 套件都包含 TensorBoard）。解除安裝單個 pip 套件可能會導致 tensorboard 可執行檔被移除，然後很難替換，因此可能有必要解除安裝所有內容並重新安裝單個版本

pip uninstall tensorflow tf-nightly tensorboard tb-nightly
pip install tensorflow

Nsight#

NVIDIA 的 Nsight 工具可用於追蹤和分析 GPU 上 JAX 程式碼的效能。有關詳細資訊，請參閱 Nsight 說明文件。