注意
点击 这里 下载完整示例代码
分析你的 PyTorch 模块¶
创建日期: 2020年12月30日 | 最后更新日期: 2024年1月19日 | 最后验证日期: 2024年11月5日
作者： 苏拉吉·苏布拉马尼安
PyTorch 包含一个性能分析器 API，有助于识别代码中各种 PyTorch 操作的时间和内存开销。性能分析器可以轻松集成到代码中，并且结果可以以表格形式打印或返回到 JSON 跟踪文件中。
注意
Profiler 支持多线程模型。Profiler 在执行操作的同一线程中运行，但它也会对可能在另一个线程中运行的子操作进行分析。同时运行的Profiler 将被限定在其各自的线程内，以防止结果混杂。
注意
PyTorch 1.8 引入了将在未来版本中取代旧的分析器 API 的新 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会带来一定的开销，最好仅用于代码分析。如果您正在测试运行时间，请记得移除它。
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
剖析前向传播¶
我们初始化随机输入张量和掩码张量，并构建模型。
在我们运行性能分析器之前，我们需要预热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.key_averages
按操作符名称聚合结果，并可选地按输入形状和/或堆栈跟踪事件。
按输入形状分组有助于识别模型使用了哪些张量形状。
这里，我们使用 group_by_stack_n=5 按操作及其回溯（截断到最近的5个事件）聚合运行时，并按注册顺序显示事件。通过传递 sort_by 参数，表格也可以进行排序（请参阅文档以获取有效的排序键）。
注意
在笔记本中运行性能分析器时，您可能会看到堆栈跟踪中出现类似<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是非常昂贵的！
aten::copy_ 运算符在 forward (12) 中复制 mask 到CPU，
以便它可以使用NumPy的 argwhere 函数。aten::copy_ 在 forward(13)
中将数组作为张量复制回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
"""