注意
跳至末尾 下載完整的示例程式碼。
分析您的 PyTorch 模組#
創建於：2020 年 12 月 30 日 | 最後更新：2024 年 1 月 19 日 | 最後驗證：2024 年 11 月 05 日
PyTorch 包含一個Profiler API，有助於識別程式碼中各種 PyTorch 操作的時間和記憶體成本。Profiler 可以輕鬆整合到您的程式碼中，並且結果可以列印為表格或以 JSON 跟蹤檔案形式返回。
注意
Profiler 支援多執行緒模型。Profiler 在與操作相同的執行緒中執行，但它也會分析可能在另一個執行緒中執行的子操作。併發執行的 Profiler 將被限制在其自己的執行緒中，以防止結果混淆。
注意
PyTorch 1.8 引入了新的 API，將在未來的版本中取代舊的 Profiler API。請檢視 此頁面 上的新 API。
請參閱 此食譜，以更快地瞭解 Profiler API 的用法。
import torch
import numpy as np
from torch import nn
import torch.autograd.profiler as profiler
使用 Profiler 進行效能除錯#
Profiler 可用於識別模型中的效能瓶頸。在此示例中，我們構建了一個執行兩個子任務的自定義模組：
對輸入進行線性變換，以及
使用變換結果在掩碼張量上獲取索引。
我們使用 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
分析前向傳播#
我們初始化隨機輸入和掩碼張量以及模型。
在執行 Profiler 之前，我們進行 CUDA 預熱，以確保準確的效能基準測試。我們將模組的前向傳播包裝在 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 引數對錶格進行排序（請參閱 文件 以獲取有效的排序鍵）。
注意
在筆記本中執行 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 的更多資訊：
