官网链接
Profiling your PyTorch Module --- PyTorch Tutorials 2.0.1+cu117 documentation
分析pytorch模块
PyTorch包含一个分析器API，用于识别代码中各种PyTorch操作的时间和内存成本。分析器可以很容易地集成到代码中，结果可以作为表格打印或以JSON跟踪文件返回。
分析器支持多线程模型。分析器与主线程在同一个线程中运行，但它也会分析可能在另一个线程中运行的子线程。同时运行的分析器将被限制在它们自己的线程中，以防止混合结果。
PyTorch 1.8引入了新的API，将在未来的版本中取代旧的分析器API。在这个页面查看新的API(this page)。
在这个攻略中(this recipe)，你可以更快速地了解Profiler API的用法。
import torch
import numpy as np
from torch import nn
import torch.autograd.profiler as profiler
使用分析器进行性能调试
分析器可以帮助识别模型中的性能瓶颈。在这个例子中，我们构建了一个自定义模块，它执行两个子任务:
- 对输入进行线性变换
- 使用转换结果来获得掩码张量上的索引。
我们使用**profiler.record_function("label")**将每个子任务的代码包装在单独的带标签的上下文管理器中。在分析器的输出中，子任务中所有操作的综合性能指标将显示在相应的标签下。
注意，使用分析器会带来一些开销，所以最好只在检查代码时使用。如果您要对运行时进行基准测试，请记住删除它。
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
分析forward执行
我们初始化随机输入和掩码张量以及模型。
在我们运行分析器之前，我们热身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.key_averages 按操作符名称、输入形状和/或堆栈跟踪事件聚合结果。按输入形状分组对于识别模型使用哪些张量形状是有用的。
在这里，我们使用group_by_stack_n=5 ，它根据操作及其回溯(截断为最近的5个事件)汇总运行时间，并按注册的顺序显示事件。也可以通过传递sort_by 参数对表进行排序(有关有效的排序键，请参阅文档docs)。
在notebook上运行分析器时，你可能会在堆栈跟踪中看到类似**<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)操作。让我们首先尝试解决内存消耗问题。可以看到，第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，这样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
"""
扩展阅读
我们已经看到了如何在PyTorch模型中使用分析器来调查时间和内存瓶颈。在这里阅读更多关于Profiler的信息: