注意
点击 这里 下载完整的示例代码
分析你的PyTorch模块¶
创建于：2020年12月30日 | 最后更新：2024年1月19日 | 最后验证：2024年11月5日
PyTorch 包含一个分析器 API，可用于识别代码中各种 PyTorch 操作的时间和内存成本。分析器可以轻松集成到您的代码中，结果可以以表格形式打印或以 JSON 跟踪文件形式返回。
注意
分析器支持多线程模型。分析器在与操作相同的线程中运行，但它也会分析可能在另一个线程中运行的子操作符。并发运行的分析器将被限定在自己的线程中，以防止结果混淆。
注意
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.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.key_averages
按操作符名称聚合结果，并可选择按输入形状和/或堆栈跟踪事件进行聚合。
按输入形状分组有助于识别模型使用了哪些张量形状。
在这里，我们使用group_by_stack_n=5，它通过操作及其回溯（截断到最近的5个事件）来聚合运行时间，并按注册顺序显示事件。表格也可以通过传递sort_by参数进行排序（请参阅文档以获取有效的排序键）。
注意
在笔记本中运行分析器时，您可能会在堆栈跟踪中看到像
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
"""
