备注
点击此处下载完整示例代码
PyTorch 模块的配置 ¶
创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日
作者：Suraj Subramanian
PyTorch 包含一个非常有用的分析器 API，可以帮助您识别代码中各种 PyTorch 操作的时间和内存成本。分析器可以轻松集成到您的代码中，结果可以打印成表格或以 JSON 追踪文件的形式返回。
备注
分析器支持多线程模型。分析器在操作所在的线程中运行，但它也会分析可能运行在另一个线程中的子操作。同时运行的多个分析器将各自限定在自己的线程中，以防止结果混合。
备注
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 可以帮助您识别模型中的性能瓶颈。在这个例子中，我们构建了一个自定义模块，该模块执行两个子任务：
对输入进行线性变换，
使用变换结果在掩码张量上获取索引。
我们使用单独的标记上下文管理器将每个子任务的代码封装起来，使用 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。如果我们在这里使用 nonzero() 的 torch 函数，就可以消除这两个步骤。
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 的更多信息，请参阅此处：
脚本总运行时间：（0 分钟 0.000 秒）
