注意
跳转至页面底部 下载完整示例代码。
Profiling your PyTorch Module (分析 PyTorch 模块)#
创建日期: 2020年12月30日 | 最后更新: 2025年11月11日 | 最后验证: 2024年11月05日
PyTorch 包含一个性能分析器（Profiler）API，可用于识别代码中各种 PyTorch 操作的时间和内存开销。Profiler 可以轻松集成到您的代码中，结果既可以以表格形式打印，也可以保存为 JSON 追踪文件。
注意
Profiler 支持多线程模型。它会在操作所在的线程中运行，同时也会分析在其他线程中运行的子算子。同时运行的多个分析器将限定在各自的线程范围内，以防止结果混淆。
注意
PyTorch 1.8 引入了全新的 API，该 API 将在未来的版本中取代旧的分析器 API。请查看 此页面 了解新 API。
前往 此方案说明 获取关于 Profiler API 使用的快速入门指南。
import torch
import numpy as np
from torch import nn
import torch.autograd.profiler as profiler
使用 Profiler 进行性能调试#
Profiler 可用于识别模型中的性能瓶颈。在本示例中，我们构建一个执行两个子任务的自定义模块：
对输入进行线性变换，以及
使用变换结果在掩码张量（mask tensor）上获取索引。
我们使用 profiler.record_function("label") 将每个子任务的代码包装在单独的带标签上下文管理器中。在分析器输出中，子任务中所有操作的汇总性能指标将显示在其对应的标签下。
请注意，使用 Profiler 会产生一定的开销，因此最好仅在调研代码时使用。如果您正在进行运行时的基准测试（benchmarking），请记得将其移除。
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 Pass）#
我们初始化随机输入、掩码张量以及模型。
在运行分析器之前，我们先预热 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 参数进行排序（有关有效的排序键，请参阅 文档）。
注意
在 Notebook 中运行分析器时，您可能会在堆栈追踪中看到类似 <ipython-input-18-193a910735e8>(13): forward 的条目，而不是文件名。这些对应于 <notebook-cell>(行号): 调用函数名。
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 的更多信息