@ariG23498: 有两篇官方PyTorch教程链接到了"Profiling in PyTorch"系列！https://docs.pytorch.org/tutoria…

有两个官方torch教程链接到“Profiling in PyTorch“系列！https://docs.pytorch.org/tutorials/beginner/profiler.html… https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html… — # 剖析你的PyTorch模块 来源：https://docs.pytorch.org/tutorials/beginner/profiler.html 注意 前往文末（https://docs.pytorch.org/tutorials/beginner/profiler.html#sphx-glr-download-beginner-profiler-py）下载完整示例代码。 创建于：2020年12月30日 | 最后更新：2026年7月19日 | 最后验证：2024年11月5日 作者： Suraj Subramanian (https://github.com/subramen) PyTorch包含一个分析器API，可用于识别代码中各种PyTorch操作的时间和内存成本。分析器可以轻松集成到你的代码中，结果可以以表格形式打印或以JSON跟踪文件形式返回。 注意 分析器支持多线程模型。分析器在与操作相同的线程中运行，但也会分析可能在另一个线程中运行的子操作符。并发运行的分析器将限定在自己的线程中，以防止结果混淆。 注意 PyTorch 1.8引入了新的API，将在未来版本中取代旧的分析器API。请查看此页面上的新API（https://pytorch.org/docs/master/profiler.html）。前往此配方（https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html）快速了解分析器API的使用。 — import torch import numpy as np from torch import nn import torch.autograd.profiler as profiler ## 使用分析器进行性能调试# (https://docs.pytorch.org/tutorials/beginner/profiler.html#performance-debugging-using-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 ## 分析前向传播# (https://docs.pytorch.org/tutorials/beginner/profiler.html#profile-the-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() # 预热 model(input, mask) with profiler.profile(with_stack=True, profile_memory=True) as prof: out, idx = model(input, mask) ## 打印分析器结果# (https://docs.pytorch.org/tutorials/beginner/profiler.html#print-profiler-results) 最后，我们打印分析器结果。profiler.key_averages按操作符名称聚合结果，并可选地按输入形状和/或堆栈跟踪事件进行聚合。按输入形状分组有助于识别模型使用了哪些张量形状。这里，我们使用group_by_stack_n=5，它按操作及其回溯（截断为最近的5个事件）聚合运行时间，并按注册顺序显示事件。表格也可以通过传递sort_by参数进行排序（有关有效的排序键，请参阅文档（https://pytorch.org/docs/stable/autograd.html#profiler））。 注意 在笔记本中运行分析器时，你可能会在堆栈跟踪中看到像(13): forward这样的条目而不是文件名。这些对应的是(行号): 调用函数。 print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5)) """ (部分列已省略) ------------- ------------ ------------ ------------ --------------------------------- 名称 Self CPU % Self CPU Self CPU Mem 源位置 ------------- ------------ ------------ ------------ --------------------------------- MASK INDICES 87.88% 5.212s -953.67 Mb /mnt/xarfuse/.../torch/au (10): forward /mnt/xarfuse/.../torch/nn (9): /mnt/xarfuse/.../IPython/ aten::copy_ 12.07% 715.848ms 0 b (12): forward /mnt/xarfuse/.../torch/nn (9): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ LINEAR PASS 0.01% 350.151us -20 b /mnt/xarfuse/.../torch/au (7): forward /mnt/xarfuse/.../torch/nn (9): /mnt/xarfuse/.../IPython/ aten::addmm 0.00% 293.342us 0 b /mnt/xarfuse/.../torch/nn /mnt/xarfuse/.../torch/nn /mnt/xarfuse/.../torch/nn (8): forward /mnt/xarfuse/.../torch/nn aten::mean 0.00% 235.095us 0 b (11): forward /mnt/xarfuse/.../torch/nn (9): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ ----------------------------- ------------ ---------- ---------------------------------- Self CPU time total: 5.931s """ ## 改进内存性能# (https://docs.pytorch.org/tutorials/beginner/profiler.html#improve-memory-performance) 注意，在内存和时间方面最昂贵的操作是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() # 预热 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)) """ (部分列已省略) ----------------- ------------ ------------ ------------ -------------------------------- 名称 Self CPU % Self CPU Self CPU Mem 源位置 ----------------- ------------ ------------ ------------ -------------------------------- MASK INDICES 93.61% 5.006s -476.84 Mb /mnt/xarfuse/.../torch/au (10): forward /mnt/xarfuse/ /torch/nn (9): /mnt/xarfuse/.../IPython/ aten::copy_ 6.34% 338.759ms 0 b (12): forward /mnt/xarfuse/.../torch/nn (9): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ aten::as_strided 0.01% 281.808us 0 b (11): forward /mnt/xarfuse/.../torch/nn (9): /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 (8): forward /mnt/xarfuse/.../torch/nn aten::_local 0.01% 268.650us 0 b (11): forward _scalar_dense /mnt/xarfuse/.../torch/nn (9): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ ----------------- ------------ ------------ ------------ -------------------------------- Self CPU time total: 5.347s """ 该操作的CPU内存占用已减半。 ## 改进时间性能# (https://docs.pytorch.org/tutorials/beginner/profiler.html#improve-time-performance) 虽然时间消耗也有所减少，但仍然过高。事实证明，将矩阵从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() # 预热 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)) """ (部分列已省略) -------------- ------------ ------------ ------------ --------------------------------- 名称 Self CPU % Self CPU Self CPU Mem 源位置 -------------- ------------ ------------ ------------ --------------------------------- aten::gt 57.17% 129.089ms 0 b (12): forward /mnt/xarfuse/.../torch/nn (25): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ aten::nonzero 37.38% 84.402ms 0 b (12): forward /mnt/xarfuse/.../torch/nn (25): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ INDEX SCORE 3.32% 7.491ms -119.21 Mb /mnt/xarfuse/.../torch/au (10): forward /mnt/xarfuse/.../torch/nn (25): /mnt/xarfuse/.../IPython/ aten::as_strided 0.20% 441.587us 0 b (12): forward /mnt/xarfuse/.../torch/nn (25): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ aten::nonzero _numpy 0.18% 395.602us 0 b (12): forward /mnt/xarfuse/.../torch/nn (25): /mnt/xarfuse/.../IPython/ /mnt/xarfuse/.../IPython/ -------------- ------------ ------------ ------------ --------------------------------- Self CPU time total: 225.801ms """ ## 延伸阅读# (https://docs.pytorch.org/tutorials/beginner/profiler.html#further-reading) 我们已经了解了如何使用分析器调查PyTorch模型中的时间和内存瓶颈。在此处阅读更多关于分析器的信息： - 分析器使用配方 (https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) - 分析基于RPC的工作负载 (https://pytorch.org/tutorials/recipes/distributed_rpc_profiling.html) - 分析器API文档 (https://pytorch.org/docs/stable/autograd.html?highlight=profiler#profiler) - Hugging Face PyTorch分析器博客文章 (https://huggingface.co/blog/torch-profiler) > Aritra 🤗 (@ariG23498): > 分析博客文章现已列入官方torch分析器文档！🤗