分析您的 PyTorch 模块
PyTorch 提供了一个性能分析器 API，可用于识别代码中各种 PyTorch 操作的时间和内存开销。性能分析器可以轻松集成到您的代码中，结果可以以表格形式打印或以 JSON 跟踪文件的形式返回。
Profiler 支持多线程模型。Profiler 在与操作相同的线程中运行，但它也会分析可能在另一个线程中运行的子操作。并发运行的 Profiler 将被限制在各自的线程中，以防止结果混淆。
PyTorch 1.8 引入了新的 API，该 API 将在未来的版本中取代旧的性能分析器 API。请在此页面查看新 API。
前往 此教程 快速了解 Profiler API 的使用方法。
importtorch
importnumpyasnp
fromtorchimport nn
importtorch.autograd.profilerasprofiler
使用 Profiler 进行性能调试
性能分析器可以帮助您识别模型中的性能瓶颈。在本例中，我们将构建一个执行两个子任务的自定义模块：
-
对输入进行线性变换，
-
使用变换结果在掩码张量上获取索引。
我们将每个子任务的代码分别包装在使用 profiler.record_function("label") 标记的上下文管理器中。在分析器的输出中，子任务中所有操作的聚合性能指标将显示在其对应的标签下。
需要注意的是，使用分析器会带来一些开销，因此最好仅用于代码的调试和性能分析。如果是在进行运行时间基准测试，请记得将其移除。
classMyModule(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)
defforward(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>(行号): 调用函数。
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() 函数，就可以消除这两个操作。
classMyModule(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)
defforward(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 的更多信息，请阅读此处：