注意
跳至末尾 下载完整的示例代码。
分析您的 PyTorch 模块#
创建于：2020 年 12 月 30 日 | 最后更新：2024 年 1 月 19 日 | 最后验证：2024 年 11 月 05 日
PyTorch 包含一个Profiler API，有助于识别代码中各种 PyTorch 操作的时间和内存成本。Profiler 可以轻松集成到您的代码中，并且结果可以打印为表格或以 JSON 跟踪文件形式返回。
注意
Profiler 支持多线程模型。Profiler 在与操作相同的线程中运行，但它也会分析可能在另一个线程中运行的子操作。并发运行的 Profiler 将被限制在其自己的线程中，以防止结果混淆。
注意
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 可用于识别模型中的性能瓶颈。在此示例中，我们构建了一个执行两个子任务的自定义模块：
对输入进行线性变换，以及
使用变换结果在掩码张量上获取索引。
我们使用 profiler.record_function("label") 将每个子任务的代码包装在单独的标记上下文管理器中。在 Profiler 输出中，子任务中所有操作的聚合性能指标将显示在其相应的标签下。
请注意，使用 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
分析前向传播#
我们初始化随机输入和掩码张量以及模型。
在运行 Profiler 之前，我们进行 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 结果#
最后，我们打印 Profiler 结果。 profiler.key_averages 按操作名称以及可选的输入形状和/或堆栈跟踪事件进行聚合结果。按输入形状分组有助于识别模型使用的张量形状。
在这里，我们使用 group_by_stack_n=5，它按操作及其回溯（截断为最近的 5 个事件）聚合运行时，并按事件注册的顺序显示它们。还可以通过传递 sort_by 参数对表格进行排序（请参阅 文档 以获取有效的排序键）。
注意
在笔记本中运行 Profiler 时，您可能会看到类似于 <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
"""
进一步阅读#
我们已经了解了 Profiler 如何用于调查 PyTorch 模型中的时间和内存瓶颈。在此处阅读有关 Profiler 的更多信息：
