注意
跳转到结尾 下载完整的示例代码。
分析你的 PyTorch 模块#
创建于：2020 年 12 月 30 日 | 最后更新：2025 年 11 月 11 日 | 最后验证：2024 年 11 月 05 日
PyTorch 包含一个分析器 API，可用于识别代码中各种 PyTorch 操作的时间和内存成本。分析器可以轻松集成到你的代码中，结果可以打印为表格或返回到 JSON 跟踪文件中。
注意
分析器支持多线程模型。分析器在与操作相同的线程中运行，但它也会分析可能在另一个线程中运行的子运算符。并发运行的分析器将限定在其自己的线程中，以防止结果混合。
注意
PyTorch 1.8 引入了将在未来版本中取代旧分析器 API 的新 API。请在此处 查看 新 API。
前往 此配方，以快速了解分析器 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
分析前向传递#
我们初始化随机输入和掩码张量，以及模型。
在运行分析器之前，我们预热 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 参数对表格进行排序（请参阅 文档 以获取有效的排序键）。
注意
在笔记本中运行分析器时，你可能会看到类似于 <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
"""
进一步阅读#
我们已经看到如何使用分析器来调查 PyTorch 模型中的时间和内存瓶颈。在此处了解有关分析器的更多信息
