注意

點選此處下載完整的示例程式碼

效能分析你的 PyTorch Module¶

建立時間：2020 年 12 月 30 日 | 最後更新：2024 年 1 月 19 日 | 最後驗證：2024 年 11 月 5 日

作者： Suraj Subramanian

PyTorch 包含一個 profiler（效能分析器）API，可用於識別程式碼中各種 PyTorch 操作的時間和記憶體開銷。Profiler 可以輕鬆地整合到你的程式碼中，並且結果可以列印為表格或以 JSON 跟蹤檔案的形式返回。

注意

Profiler 支援多執行緒模型。Profiler 在與操作相同的執行緒中執行，但它也會對可能在另一個執行緒中執行的子運算元進行效能分析。併發執行的 profiler 將限定在其各自的執行緒範圍內，以防止結果混雜。

注意

PyTorch 1.8 引入了新的 API，該 API 將在未來版本中取代舊的 profiler API。請訪問此頁面檢視新 API。

前往這個程式碼示例，以便更快地瞭解 Profiler API 的用法。

import torch
import numpy as np
from torch import nn
import torch.autograd.profiler as profiler

使用 Profiler 進行效能除錯¶

Profiler 可用於識別模型中的效能瓶頸。在本示例中，我們構建了一個執行兩個子任務的自定義 module：

對輸入進行線性變換，以及
使用變換結果獲取 mask 張量上的索引。

我們使用 profiler.record_function("label") 將每個子任務的程式碼包裝在單獨的帶標籤的上下文管理器中。在 profiler 輸出中，子任務中所有操作的聚合效能指標將顯示在其對應的標籤下。

請注意，使用 Profiler 會產生一定的開銷，最好僅用於程式碼調查。如果你正在進行執行時基準測試，請記住將其移除。

class MyModule(nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super(MyModule, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias)

    def forward(self, input, mask):
        with profiler.record_function("LINEAR PASS"):
            out = self.linear(input)

        with profiler.record_function("MASK INDICES"):
            threshold = out.sum(axis=1).mean().item()
            hi_idx = np.argwhere(mask.cpu().numpy() > threshold)
            hi_idx = torch.from_numpy(hi_idx).cuda()

        return out, hi_idx

對前向傳播進行效能分析¶

我們初始化隨機的輸入張量和 mask 張量，以及模型。

在執行 profiler 之前，我們先對 CUDA 進行預熱，以確保精確的效能基準測試。我們將 module 的前向傳播包裝在 profiler.profile 上下文管理器中。with_stack=True 引數會在跟蹤中附加操作的檔名和行號。

警告

with_stack=True 會產生額外的開銷，更適合用於程式碼調查。如果你正在進行效能基準測試，請記住將其移除。

model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.double).cuda()

# warm-up
model(input, mask)

with profiler.profile(with_stack=True, profile_memory=True) as prof:
    out, idx = model(input, mask)

列印 profiler 結果¶

最後，我們列印 profiler 結果。profiler.key_averages 按運算元名稱聚合結果，也可以選擇按輸入形狀和/或堆疊跟蹤事件進行聚合。按輸入形狀分組有助於識別模型使用了哪些張量形狀。

在這裡，我們使用 group_by_stack_n=5，它按操作及其回溯（截斷到最近的 5 個事件）聚合執行時，並按事件註冊的順序顯示。表格也可以透過傳遞 sort_by 引數進行排序（請參考文件瞭解有效的排序鍵）。

注意

在 notebook 中執行 profiler 時，你在堆疊跟蹤中可能會看到諸如 <ipython-input-18-193a910735e8>(13): forward 的條目，而不是檔名。這些對應於 <notebook-cell>(line number): calling-function。

print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5))

"""
(Some columns are omitted)

-------------  ------------  ------------  ------------  ---------------------------------
         Name    Self CPU %      Self CPU  Self CPU Mem   Source Location
-------------  ------------  ------------  ------------  ---------------------------------
 MASK INDICES        87.88%        5.212s    -953.67 Mb  /mnt/xarfuse/.../torch/au
                                                         <ipython-input-...>(10): forward
                                                         /mnt/xarfuse/.../torch/nn
                                                         <ipython-input-...>(9): <module>
                                                         /mnt/xarfuse/.../IPython/

  aten::copy_        12.07%     715.848ms           0 b  <ipython-input-...>(12): forward
                                                         /mnt/xarfuse/.../torch/nn
                                                         <ipython-input-...>(9): <module>
                                                         /mnt/xarfuse/.../IPython/
                                                         /mnt/xarfuse/.../IPython/

  LINEAR PASS         0.01%     350.151us         -20 b  /mnt/xarfuse/.../torch/au
                                                         <ipython-input-...>(7): forward
                                                         /mnt/xarfuse/.../torch/nn
                                                         <ipython-input-...>(9): <module>
                                                         /mnt/xarfuse/.../IPython/

  aten::addmm         0.00%     293.342us           0 b  /mnt/xarfuse/.../torch/nn
                                                         /mnt/xarfuse/.../torch/nn
                                                         /mnt/xarfuse/.../torch/nn
                                                         <ipython-input-...>(8): forward
                                                         /mnt/xarfuse/.../torch/nn

   aten::mean         0.00%     235.095us           0 b  <ipython-input-...>(11): forward
                                                         /mnt/xarfuse/.../torch/nn
                                                         <ipython-input-...>(9): <module>
                                                         /mnt/xarfuse/.../IPython/
                                                         /mnt/xarfuse/.../IPython/

-----------------------------  ------------  ---------- ----------------------------------
Self CPU time total: 5.931s

"""

改進記憶體效能¶

請注意，記憶體和時間方面開銷最大的操作位於 forward (10)，它代表 MASK INDICES 中的操作。我們先嚐試解決記憶體消耗問題。可以看到，第 12 行的 .to() 操作消耗了 953.67 Mb。此操作將 mask 複製到 CPU。mask 使用 torch.double 資料型別初始化。我們能否透過將其轉換為 torch.float 來減少記憶體佔用？

model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.float).cuda()

# warm-up
model(input, mask)

with profiler.profile(with_stack=True, profile_memory=True) as prof:
    out, idx = model(input, mask)

print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5))

"""
(Some columns are omitted)

-----------------  ------------  ------------  ------------  --------------------------------
             Name    Self CPU %      Self CPU  Self CPU Mem   Source Location
-----------------  ------------  ------------  ------------  --------------------------------
     MASK INDICES        93.61%        5.006s    -476.84 Mb  /mnt/xarfuse/.../torch/au
                                                             <ipython-input-...>(10): forward
                                                             /mnt/xarfuse/  /torch/nn
                                                             <ipython-input-...>(9): <module>
                                                             /mnt/xarfuse/.../IPython/

      aten::copy_         6.34%     338.759ms           0 b  <ipython-input-...>(12): forward
                                                             /mnt/xarfuse/.../torch/nn
                                                             <ipython-input-...>(9): <module>
                                                             /mnt/xarfuse/.../IPython/
                                                             /mnt/xarfuse/.../IPython/

 aten::as_strided         0.01%     281.808us           0 b  <ipython-input-...>(11): forward
                                                             /mnt/xarfuse/.../torch/nn
                                                             <ipython-input-...>(9): <module>
                                                             /mnt/xarfuse/.../IPython/
                                                             /mnt/xarfuse/.../IPython/

      aten::addmm         0.01%     275.721us           0 b  /mnt/xarfuse/.../torch/nn
                                                             /mnt/xarfuse/.../torch/nn
                                                             /mnt/xarfuse/.../torch/nn
                                                             <ipython-input-...>(8): forward
                                                             /mnt/xarfuse/.../torch/nn

      aten::_local        0.01%     268.650us           0 b  <ipython-input-...>(11): forward
      _scalar_dense                                          /mnt/xarfuse/.../torch/nn
                                                             <ipython-input-...>(9): <module>
                                                             /mnt/xarfuse/.../IPython/
                                                             /mnt/xarfuse/.../IPython/

-----------------  ------------  ------------  ------------  --------------------------------
Self CPU time total: 5.347s

"""

此操作的 CPU 記憶體佔用減少了一半。

改進時間效能¶

雖然消耗的時間也減少了一些，但仍然太高。事實證明，將矩陣從 CUDA 複製到 CPU 非常耗時！forward (12) 中的 aten::copy_ 運算元將 mask 複製到 CPU，以便可以使用 NumPy 的 argwhere 函式。forward(13) 中的 aten::copy_ 將陣列作為張量複製回 CUDA。如果我們在此處改用 torch 函式 nonzero()，就可以消除這兩個操作。

class MyModule(nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super(MyModule, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias)

    def forward(self, input, mask):
        with profiler.record_function("LINEAR PASS"):
            out = self.linear(input)

        with profiler.record_function("MASK INDICES"):
            threshold = out.sum(axis=1).mean()
            hi_idx = (mask > threshold).nonzero(as_tuple=True)

        return out, hi_idx


model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.float).cuda()

# warm-up
model(input, mask)

with profiler.profile(with_stack=True, profile_memory=True) as prof:
    out, idx = model(input, mask)

print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5))

"""
(Some columns are omitted)

--------------  ------------  ------------  ------------  ---------------------------------
          Name    Self CPU %      Self CPU  Self CPU Mem   Source Location
--------------  ------------  ------------  ------------  ---------------------------------
      aten::gt        57.17%     129.089ms           0 b  <ipython-input-...>(12): forward
                                                          /mnt/xarfuse/.../torch/nn
                                                          <ipython-input-...>(25): <module>
                                                          /mnt/xarfuse/.../IPython/
                                                          /mnt/xarfuse/.../IPython/

 aten::nonzero        37.38%      84.402ms           0 b  <ipython-input-...>(12): forward
                                                          /mnt/xarfuse/.../torch/nn
                                                          <ipython-input-...>(25): <module>
                                                          /mnt/xarfuse/.../IPython/
                                                          /mnt/xarfuse/.../IPython/

   INDEX SCORE         3.32%       7.491ms    -119.21 Mb  /mnt/xarfuse/.../torch/au
                                                          <ipython-input-...>(10): forward
                                                          /mnt/xarfuse/.../torch/nn
                                                          <ipython-input-...>(25): <module>
                                                          /mnt/xarfuse/.../IPython/

aten::as_strided         0.20%    441.587us          0 b  <ipython-input-...>(12): forward
                                                          /mnt/xarfuse/.../torch/nn
                                                          <ipython-input-...>(25): <module>
                                                          /mnt/xarfuse/.../IPython/
                                                          /mnt/xarfuse/.../IPython/

 aten::nonzero
     _numpy             0.18%     395.602us           0 b  <ipython-input-...>(12): forward
                                                          /mnt/xarfuse/.../torch/nn
                                                          <ipython-input-...>(25): <module>
                                                          /mnt/xarfuse/.../IPython/
                                                          /mnt/xarfuse/.../IPython/
--------------  ------------  ------------  ------------  ---------------------------------
Self CPU time total: 225.801ms

"""

延伸閱讀¶

我們已經瞭解瞭如何使用 Profiler 來調查 PyTorch 模型中的時間與記憶體瓶頸。在此處閱讀更多關於 Profiler 的資訊：

指令碼總執行時間： ( 0 分 0.000 秒)

由 Sphinx-Gallery 生成的畫廊