使用 TensorBoard 的 PyTorch 分析器

建立於：2021 年 4 月 20 日 | 上次更新：2024 年 10 月 31 日 | 上次驗證：2024 年 11 月 05 日

本教學課程示範如何將 TensorBoard 外掛程式與 PyTorch 分析器搭配使用，以偵測模型的效能瓶頸。

警告

TensorBoard 與 PyTorch 分析器的整合現在已被棄用。 請改用 Perfetto 或 Chrome 追蹤來檢視 trace.json 檔案。 在產生追蹤後，只需將 trace.json 拖曳到 Perfetto UI 或 chrome://tracing 以視覺化您的設定檔。

簡介

PyTorch 1.8 包含更新的分析器 API，能夠記錄 CPU 端的運算以及 GPU 端的 CUDA 核心啟動。分析器可以在 TensorBoard 外掛程式中視覺化此資訊，並提供效能瓶頸的分析。

在本教學課程中，我們將使用一個簡單的 Resnet 模型來示範如何使用 TensorBoard 外掛程式來分析模型效能。

設定

若要安裝 torch 和 torchvision，請使用以下指令

pip install torch torchvision

步驟

1. 準備資料和模型

2. 使用分析器記錄執行事件

3. 執行分析器

4. 使用 TensorBoard 檢視結果並分析模型效能

5. 在分析器的幫助下提高效能

6. 使用其他進階功能分析效能

7. 其他實作：在 AMD GPU 上分析 PyTorch

1. 準備資料和模型

首先，匯入所有必要的程式庫

import torch
import torch.nn
import torch.optim
import torch.profiler
import torch.utils.data
import torchvision.datasets
import torchvision.models
import torchvision.transforms as T

然後準備輸入資料。對於本教學課程，我們使用 CIFAR10 資料集。將其轉換為所需的格式，並使用 DataLoader 來載入每個批次。

transform = T.Compose(
    [T.Resize(224),
     T.ToTensor(),
     T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)

接下來，建立 Resnet 模型、損失函數和最佳化器物件。若要在 GPU 上執行，請將模型和損失移至 GPU 裝置。

device = torch.device("cuda:0")
model = torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device)
criterion = torch.nn.CrossEntropyLoss().cuda(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()

定義每個輸入資料批次的訓練步驟。

def train(data):
    inputs, labels = data[0].to(device=device), data[1].to(device=device)
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

2. 使用分析器記錄執行事件

分析器透過上下文管理器啟用，並接受多個參數，其中一些最有用的參數為

• schedule - 可呼叫物件，將 step (int) 作為單一參數，並傳回要在每個步驟執行的分析器動作。

  在本範例中，使用 wait=1, warmup=1, active=3, repeat=1，分析器將跳過第一個步驟/迭代，在第二個步驟開始熱身，記錄接下來的三個迭代，之後追蹤將可用，並且呼叫 on_trace_ready（如果已設定）。總之，該循環重複一次。每個循環在 TensorBoard 外掛程式中稱為「span」。

  在 wait 階段，profiler 會被停用。在 warmup 階段，profiler 會開始追蹤，但結果會被捨棄。這是為了減少 profiling 的額外負擔。profiling 一開始的負擔很高，容易使 profiling 結果產生偏差。在 active 階段，profiler 會運作並記錄事件。

• on_trace_ready - 可呼叫的物件，會在每個週期的結束時被呼叫；在這個範例中，我們使用 torch.profiler.tensorboard_trace_handler 來產生 TensorBoard 的結果檔案。在 profiling 之後，結果檔案會被儲存到 ./log/resnet18 目錄中。將這個目錄指定為 logdir 參數，以便在 TensorBoard 中分析 profile。

• record_shapes - 是否記錄運算子輸入的形狀。

• profile_memory - 追蹤 tensor 的記憶體配置/釋放。請注意，對於 1.10 之前的舊版 PyTorch，如果您遇到長時間的 profiling，請停用它或升級到新版本。

• with_stack - 記錄運算子的原始碼資訊（檔案和行號）。如果 TensorBoard 是在 VS Code 中啟動的（參考），點擊堆疊框架會導航到特定的程式碼行。

with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
        record_shapes=True,
        profile_memory=True,
        with_stack=True
) as prof:
    for step, batch_data in enumerate(train_loader):
        prof.step()  # Need to call this at each step to notify profiler of steps' boundary.
        if step >= 1 + 1 + 3:
            break
        train(batch_data)

或者，也支援以下非上下文管理器的啟動/停止方式。

prof = torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
        record_shapes=True,
        with_stack=True)
prof.start()
for step, batch_data in enumerate(train_loader):
    prof.step()
    if step >= 1 + 1 + 3:
        break
    train(batch_data)
prof.stop()

3. 執行 Profiler

執行上面的程式碼。profiling 結果將會儲存在 ./log/resnet18 目錄下。

4. 使用 TensorBoard 檢視結果並分析模型效能

注意

TensorBoard Plugin 的支援已被棄用，因此某些功能可能無法像以前一樣運作。請查看替代方案 HTA。

安裝 PyTorch Profiler TensorBoard Plugin。

pip install torch_tb_profiler

啟動 TensorBoard。

tensorboard --logdir=./log

在 Google Chrome 瀏覽器或 Microsoft Edge 瀏覽器中開啟 TensorBoard profile URL (不支援 Safari)。

http://localhost:6006/#pytorch_profiler

您應該可以看到如下所示的 Profiler plugin 頁面。

• 總覽 (Overview)

總覽顯示模型效能的高階摘要。

「GPU Summary」面板顯示 GPU 配置、GPU 使用率和 Tensor Cores 使用率。在這個範例中，GPU 使用率很低。這些指標的詳細資訊請參考這裡。

「Step Time Breakdown」顯示每個步驟在不同執行類別中所花費的時間分佈。在這個範例中，您可以看到 DataLoader 的額外負擔非常顯著。

底部的「Performance Recommendation」使用 profiling 資料自動突顯可能的瓶頸，並提供可操作的優化建議。

您可以在左側的「Views」下拉式選單中更改檢視頁面。

• 運算子檢視 (Operator view)

運算子檢視顯示每個在主機或裝置上執行的 PyTorch 運算子的效能。

「Self」持續時間不包含其子運算子的時間。「Total」持續時間包含其子運算子的時間。

• 檢視呼叫堆疊 (View call stack)

點擊運算子的 View Callstack，將會顯示具有相同名稱但不同呼叫堆疊的運算子。然後點擊此子表格中的 View Callstack，將會顯示呼叫堆疊框架。

如果 TensorBoard 是在 VS Code 中啟動的（Launch Guide），點擊呼叫堆疊框架會導航到特定的程式碼行。

• 核心檢視 (Kernel view)

GPU 核心檢視顯示所有核心在 GPU 上所花費的時間。

Tensor Cores Used：這個核心是否使用 Tensor Cores。

Mean Blocks per SM：Blocks per SM = 這個核心的 Blocks / 這個 GPU 的 SM 數量。如果這個數字小於 1，則表示 GPU 多處理器未被完全利用。「Mean Blocks per SM」是這個核心名稱的所有運行的加權平均值，使用每次運行的持續時間作為權重。

Mean Est. Achieved Occupancy：Est. Achieved Occupancy 在此欄位的工具提示中定義。對於大多數情況（例如記憶體頻寬限制的核心），數值越高越好。「Mean Est. Achieved Occupancy」是這個核心名稱的所有運行的加權平均值，使用每次運行的持續時間作為權重。

• 追蹤檢視 (Trace view)

追蹤檢視顯示已 profiling 的運算子和 GPU 核心的時間軸。您可以選擇它來查看如下所示的詳細資訊。

您可以使用右側的工具列移動圖表和放大/縮小。也可以使用鍵盤在時間軸內放大和移動。'w' 和 's' 鍵以滑鼠為中心放大，'a' 和 'd' 鍵將時間軸左右移動。您可以多次按下這些鍵，直到看到可讀的表示。

如果向後運算子的「Incoming Flow」欄位的值為「forward correspond to backward」，您可以點擊文字以取得其啟動的向前運算子。

在這個範例中，我們可以發現以 enumerate(DataLoader) 為前綴的事件花費了大量時間。並且在這段時間的大部分時間裡，GPU 都是閒置的。因為這個函數在主機端載入資料和轉換資料，在此期間 GPU 資源被浪費了。

5. 藉由 profiler 改善效能

在「Overview」頁面的底部，「Performance Recommendation」中的建議暗示瓶頸是 DataLoader。預設情況下，PyTorch DataLoader 使用單一進程。使用者可以透過設定參數 num_workers 來啟用多進程資料載入。這裡有更多詳細資訊。

在這個範例中，我們遵循「效能建議」，將 num_workers 設定如下，並將不同的名稱（例如 ./log/resnet18_4workers）傳遞給 tensorboard_trace_handler，然後再次執行。

train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=4)

接著，讓我們在左側的「Runs」下拉式選單中選擇最近分析的執行。

從上面的視圖中，我們可以發現步驟時間減少到約 76 毫秒，而之前的執行時間為 132 毫秒，其中 DataLoader 的時間減少是主要貢獻。

從上面的視圖中，我們可以看到 enumerate(DataLoader) 的執行時間減少，並且 GPU 的使用率提高。

6. 使用其他進階功能分析效能

• 記憶體視圖

若要分析記憶體，必須在 torch.profiler.profile 的引數中將 profile_memory 設定為 True。

您可以嘗試使用 Azure 上的現有範例

pip install azure-storage-blob
tensorboard --logdir=https://torchtbprofiler.blob.core.windows.net/torchtbprofiler/demo/memory_demo_1_10

分析器會記錄所有記憶體配置/釋放事件以及分析期間分配器的內部狀態。記憶體視圖由三個組件組成，如下所示。

這些組件分別是記憶體曲線圖、記憶體事件表和記憶體統計表，從上到下排列。

可以在「Device」選擇框中選擇記憶體類型。 例如，「GPU0」表示下表僅顯示每個運算子在 GPU 0 上的記憶體使用量，不包括 CPU 或其他 GPU。

記憶體曲線顯示記憶體消耗的趨勢。「Allocated」曲線顯示實際使用的總記憶體，例如張量。 在 PyTorch 中，CUDA 分配器和一些其他分配器採用快取機制。「Reserved」曲線顯示分配器保留的總記憶體。 您可以按住滑鼠左鍵並在圖表上拖曳以選擇所需範圍內的事件

選擇後，三個組件將針對受限制的時間範圍進行更新，以便您可以獲得更多相關資訊。 通過重複此過程，您可以縮放到非常細微的細節。 在圖表上按一下滑鼠右鍵將將圖表重設為初始狀態。

在記憶體事件表中，分配和釋放事件配對成一個條目。「operator」欄顯示導致分配的直接 ATen 運算子。 請注意，在 PyTorch 中，ATen 運算子通常使用 aten::empty 來分配記憶體。 例如，aten::ones 實現為 aten::empty，後跟 aten::fill_。 僅顯示運算子名稱為 aten::empty 的幫助不大。 在這種特殊情況下，它將顯示為 aten::ones (aten::empty)。「Allocation Time」、「Release Time」和「Duration」欄的資料如果事件發生在時間範圍之外，可能會遺失。

在記憶體統計表中，「Size Increase」欄總結了所有分配大小，並減去所有記憶體釋放大小，即此運算子之後記憶體使用量的淨增加。「Self Size Increase」欄與「Size Increase」類似，但不計算子運算子的分配。 關於 ATen 運算子的實現細節，一些運算子可能會調用其他運算子，因此記憶體分配可能發生在調用堆疊的任何層級。 也就是說，「Self Size Increase」僅計算調用堆疊目前層級的記憶體使用量增加。 最後，「Allocation Size」欄總結了所有分配，而不考慮記憶體釋放。

• 分散式視圖

此外掛程式現在支援使用 NCCL/GLOO 作為後端的 DDP 分析分散式視圖。

您可以嘗試使用 Azure 上的現有範例

pip install azure-storage-blob
tensorboard --logdir=https://torchtbprofiler.blob.core.windows.net/torchtbprofiler/demo/distributed_bert

「Computation/Communication Overview」顯示計算/通訊比率及其重疊程度。 從這個視圖中，使用者可以找出工作者之間的負載平衡問題。 例如，如果一個工作者的計算+重疊時間遠大於其他工作者，則可能存在負載平衡問題，或者這個工作者可能是一個落後者。

「Synchronizing/Communication Overview」顯示通訊效率。「Data Transfer Time」是實際資料交換的時間。「Synchronizing Time」是等待並與其他工作者同步的時間。

如果一個工作者的「Synchronizing Time」遠小於其他工作者，則此工作者可能是一個落後者，其計算工作量可能比其他工作者多。

「Communication Operations Stats」總結了每個工作者中所有通訊操作的詳細統計資訊。

7. 其他實作：在 AMD GPU 上分析 PyTorch

AMD ROCm 平台是一個開源軟體堆疊，專為 GPU 計算而設計，由驅動程式、開發工具和 API 組成。 我們可以在 AMD GPU 上執行上述步驟。 在本節中，我們將使用 Docker 安裝 ROCm 基本開發映像，然後安裝 PyTorch。

為了舉例說明，讓我們建立一個名為 profiler_tutorial 的目錄，並將步驟 1 中的程式碼另存為此目錄中的 test_cifar10.py。

mkdir ~/profiler_tutorial
cd profiler_tutorial
vi test_cifar10.py

在撰寫本文時，ROCm 平台上 PyTorch 的 Stable(2.1.1) Linux 版本是 ROCm 5.6。

• 從 Docker Hub 取得安裝了正確使用者空間 ROCm 版本的基礎 Docker 映像。

它是 rocm/dev-ubuntu-20.04:5.6。

• 啟動 ROCm 基礎 Docker 容器

docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 8G -v ~/profiler_tutorial:/profiler_tutorial rocm/dev-ubuntu-20.04:5.6

• 在容器內，安裝安裝 wheels 套件所需的任何依賴項。

sudo apt update
sudo apt install libjpeg-dev python3-dev -y
pip3 install wheel setuptools
sudo apt install python-is-python3

• 安裝 wheels

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6

• 安裝 torch_tb_profiler，然後執行 Python 檔案 test_cifar10.py

pip install torch_tb_profiler
cd /profiler_tutorial
python test_cifar10.py

現在，我們擁有在 TensorBoard 中檢視所需的所有資料

tensorboard --logdir=./log

如步驟 4 中所述，選擇不同的視圖。 例如，下面是 Operator 視圖

在撰寫本節時，Trace 視圖無法運作，並且不顯示任何內容。 您可以通過在 Chrome 瀏覽器中鍵入 chrome://tracing 來解決此問題。

• 將 trace.json 檔案從 ~/profiler_tutorial/log/resnet18 目錄複製到 Windows。

如果檔案位於遠端位置，您可能需要使用 scp 複製檔案。

• 按一下 Load 按鈕，從瀏覽器中的 chrome://tracing 頁面載入追蹤 JSON 檔案。

如先前所述，您可以移動圖表並放大和縮小。您也可以使用鍵盤來縮放和在時間軸內移動。w 和 s 鍵會以滑鼠游標為中心進行放大，a 和 d 鍵則會左右移動時間軸。您可以多次按下這些按鍵，直到看到可讀的呈現方式。

深入了解

請查看以下文件以繼續學習，並隨時開啟一個 issue 這裡。

• PyTorch TensorBoard Profiler Github

• torch.profiler API

• HTA