注意
跳转至页面底部 下载完整示例代码。
PyTorch Profiler#
创建日期: 2021年1月29日 | 最后更新: 2025年7月9日 | 最后验证: 未验证
本教程介绍了如何使用 PyTorch 分析器（Profiler）来衡量模型算子的时间和内存消耗。
简介#
PyTorch 包含一个简单的分析器 API，当用户需要确定模型中最耗时的算子时，该 API 非常有用。
在本教程中，我们将使用一个简单的 Resnet 模型来演示如何使用分析器来分析模型性能。
先决条件#
torch >= 2.3.0
设置#
安装 torch 和 torchvision 请使用以下命令
pip install torch torchvision
步骤#
导入所有必要的库
实例化一个简单的 Resnet 模型
使用分析器分析执行时间
使用分析器分析内存消耗
使用追踪（Tracing）功能
检查堆栈追踪
使用分析器分析长时间运行的任务
1. 导入所有必要的库#
在本教程中，我们将使用 torch、torchvision.models 和 profiler 模块。
import torch
import torchvision.models as models
from torch.profiler import profile, ProfilerActivity, record_function
2. 实例化一个简单的 Resnet 模型#
让我们创建一个 Resnet 模型实例并为其准备输入。
model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)
3. 使用分析器分析执行时间#
PyTorch 分析器通过上下文管理器启用，并接受多个参数，其中一些最有用的参数包括：
activities- 要分析的活动列表ProfilerActivity.CPU- PyTorch 算子、TorchScript 函数和用户定义的代码标签（参见下文的record_function）；ProfilerActivity.CUDA- 设备上的 CUDA 内核；ProfilerActivity.XPU- 设备上的 XPU 内核；
record_shapes- 是否记录算子输入的形状；profile_memory- 是否报告模型张量消耗的内存量；
注意：使用 CUDA 时，分析器还会显示主机上发生的运行时 CUDA 事件。
让我们看看如何使用分析器来分析执行时间。
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
with record_function("model_inference"):
model(inputs)
请注意，我们可以使用 record_function 上下文管理器来标记任意代码范围，并使用用户提供的名称（在上面的示例中，model_inference 被用作标签）。
分析器允许用户检查在分析器上下文管理器包装的代码范围执行期间调用了哪些算子。如果同时有多个分析器范围处于活动状态（例如在并行的 PyTorch 线程中），每个分析上下文管理器仅追踪其对应范围内的算子。分析器还会自动对通过 torch.jit._fork 启动的异步任务以及（在反向传播的情况下）通过 backward() 调用启动的反向传播算子进行分析。
让我们打印出上述执行的统计信息。
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
输出如下所示（省略了部分列）：
--------------------------------- ------------ ------------ ------------ ------------
Name Self CPU CPU total CPU time avg # of Calls
--------------------------------- ------------ ------------ ------------ ------------
model_inference 5.509ms 57.503ms 57.503ms 1
aten::conv2d 231.000us 31.931ms 1.597ms 20
aten::convolution 250.000us 31.700ms 1.585ms 20
aten::_convolution 336.000us 31.450ms 1.573ms 20
aten::mkldnn_convolution 30.838ms 31.114ms 1.556ms 20
aten::batch_norm 211.000us 14.693ms 734.650us 20
aten::_batch_norm_impl_index 319.000us 14.482ms 724.100us 20
aten::native_batch_norm 9.229ms 14.109ms 705.450us 20
aten::mean 332.000us 2.631ms 125.286us 21
aten::select 1.668ms 2.292ms 8.988us 255
--------------------------------- ------------ ------------ ------------ ------------
Self CPU time total: 57.549m
这里我们可以看到，正如预期的那样，大部分时间都花在卷积上（对于编译时支持 MKL-DNN 的 PyTorch，特别是 mkldnn_convolution）。注意“自 CPU 时间”（self cpu time）和“CPU 时间”（cpu time）之间的区别——算子可以调用其他算子，“自 CPU 时间”不包括在子算子调用中花费的时间，而“总 CPU 时间”包括了这些时间。你可以通过向 table 调用传递 sort_by="self_cpu_time_total" 来选择按“自 CPU 时间”排序。
为了获得更细粒度的结果并包含算子输入形状，请传递 group_by_input_shape=True（注意：这要求在运行分析器时设置 record_shapes=True）。
print(
prof.key_averages(group_by_input_shape=True).table(
sort_by="cpu_time_total", row_limit=10
)
)
输出可能如下所示（省略了部分列）：
--------------------------------- ------------ -------------------------------------------
Name CPU total Input Shapes
--------------------------------- ------------ -------------------------------------------
model_inference 57.503ms []
aten::conv2d 8.008ms [5,64,56,56], [64,64,3,3], [], ..., []]
aten::convolution 7.956ms [[5,64,56,56], [64,64,3,3], [], ..., []]
aten::_convolution 7.909ms [[5,64,56,56], [64,64,3,3], [], ..., []]
aten::mkldnn_convolution 7.834ms [[5,64,56,56], [64,64,3,3], [], ..., []]
aten::conv2d 6.332ms [[5,512,7,7], [512,512,3,3], [], ..., []]
aten::convolution 6.303ms [[5,512,7,7], [512,512,3,3], [], ..., []]
aten::_convolution 6.273ms [[5,512,7,7], [512,512,3,3], [], ..., []]
aten::mkldnn_convolution 6.233ms [[5,512,7,7], [512,512,3,3], [], ..., []]
aten::conv2d 4.751ms [[5,256,14,14], [256,256,3,3], [], ..., []]
--------------------------------- ------------ -------------------------------------------
Self CPU time total: 57.549ms
注意 aten::convolution 出现了两次，且输入形状不同。
分析器也可以用来分析在 GPU 上执行的模型的性能：用户可以在 cpu、cuda 和 xpu 之间切换。
activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
device = "cuda"
activities += [ProfilerActivity.CUDA]
elif torch.xpu.is_available():
device = "xpu"
activities += [ProfilerActivity.XPU]
else:
print(
"Neither CUDA nor XPU devices are available to demonstrate profiling on acceleration devices"
)
import sys
sys.exit(0)
sort_by_keyword = device + "_time_total"
model = models.resnet18().to(device)
inputs = torch.randn(5, 3, 224, 224).to(device)
with profile(activities=activities, record_shapes=True) as prof:
with record_function("model_inference"):
model(inputs)
print(prof.key_averages().table(sort_by=sort_by_keyword, row_limit=10))
（注意：首次使用 CUDA 分析可能会带来额外的开销。）
最终的表格输出（省略了部分列）：
------------------------------------------------------- ------------ ------------
Name Self CUDA CUDA total
------------------------------------------------------- ------------ ------------
model_inference 0.000us 11.666ms
aten::conv2d 0.000us 10.484ms
aten::convolution 0.000us 10.484ms
aten::_convolution 0.000us 10.484ms
aten::_convolution_nogroup 0.000us 10.484ms
aten::thnn_conv2d 0.000us 10.484ms
aten::thnn_conv2d_forward 10.484ms 10.484ms
void at::native::im2col_kernel<float>(long, float co... 3.844ms 3.844ms
sgemm_32x32x32_NN 3.206ms 3.206ms
sgemm_32x32x32_NN_vec 3.093ms 3.093ms
------------------------------------------------------- ------------ ------------
Self CPU time total: 23.015ms
Self CUDA time total: 11.666ms
（注意：首次使用 XPU 分析可能会带来额外的开销。）
最终的表格输出（省略了部分列）：
------------------------------ ------------ ------------ ------------ ------------ ------------
Name Self XPU Self XPU % XPU total XPU time avg # of Calls
------------------------------ ------------ ------------ ------------ ------------ ------------
model_inference 0.000us 0.00% 2.567ms 2.567ms 1
aten::conv2d 0.000us 0.00% 1.871ms 93.560us 20
aten::convolution 0.000us 0.00% 1.871ms 93.560us 20
aten::_convolution 0.000us 0.00% 1.871ms 93.560us 20
aten::convolution_overrideable 1.871ms 72.89% 1.871ms 93.560us 20
gen_conv 1.484ms 57.82% 1.484ms 74.216us 20
aten::batch_norm 0.000us 0.00% 432.640us 21.632us 20
aten::_batch_norm_impl_index 0.000us 0.00% 432.640us 21.632us 20
aten::native_batch_norm 432.640us 16.85% 432.640us 21.632us 20
conv_reorder 386.880us 15.07% 386.880us 6.448us 60
------------------------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 712.486ms
Self XPU time total: 2.567ms
注意输出中出现了设备上的内核（例如 CUDA 的 sgemm_32x32x32_NN 或 XPU 的 gen_conv）。
4. 使用分析器分析内存消耗#
PyTorch 分析器还可以显示在模型算子执行期间分配（或释放）的内存量（由模型张量使用）。在下面的输出中，“self”内存对应于算子分配（或释放）的内存，不包括对其他算子的子调用。要启用内存分析功能，请传递 profile_memory=True。
model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)
with profile(
activities=[ProfilerActivity.CPU], profile_memory=True, record_shapes=True
) as prof:
model(inputs)
print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))
# (omitting some columns)
# --------------------------------- ------------ ------------ ------------
# Name CPU Mem Self CPU Mem # of Calls
# --------------------------------- ------------ ------------ ------------
# aten::empty 94.79 Mb 94.79 Mb 121
# aten::max_pool2d_with_indices 11.48 Mb 11.48 Mb 1
# aten::addmm 19.53 Kb 19.53 Kb 1
# aten::empty_strided 572 b 572 b 25
# aten::resize_ 240 b 240 b 6
# aten::abs 480 b 240 b 4
# aten::add 160 b 160 b 20
# aten::masked_select 120 b 112 b 1
# aten::ne 122 b 53 b 6
# aten::eq 60 b 30 b 2
# --------------------------------- ------------ ------------ ------------
# Self CPU time total: 53.064ms
print(prof.key_averages().table(sort_by="cpu_memory_usage", row_limit=10))
输出可能如下所示（省略了部分列）：
--------------------------------- ------------ ------------ ------------
Name CPU Mem Self CPU Mem # of Calls
--------------------------------- ------------ ------------ ------------
aten::empty 94.79 Mb 94.79 Mb 121
aten::batch_norm 47.41 Mb 0 b 20
aten::_batch_norm_impl_index 47.41 Mb 0 b 20
aten::native_batch_norm 47.41 Mb 0 b 20
aten::conv2d 47.37 Mb 0 b 20
aten::convolution 47.37 Mb 0 b 20
aten::_convolution 47.37 Mb 0 b 20
aten::mkldnn_convolution 47.37 Mb 0 b 20
aten::max_pool2d 11.48 Mb 0 b 1
aten::max_pool2d_with_indices 11.48 Mb 11.48 Mb 1
--------------------------------- ------------ ------------ ------------
Self CPU time total: 53.064ms
5. 使用追踪（Tracing）功能#
分析结果可以输出为 .json 追踪文件：追踪 CUDA 或 XPU 内核，用户可以在 cpu、cuda 和 xpu 之间切换。
activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
device = "cuda"
activities += [ProfilerActivity.CUDA]
elif torch.xpu.is_available():
device = "xpu"
activities += [ProfilerActivity.XPU]
else:
print(
"Neither CUDA nor XPU devices are available to demonstrate profiling on acceleration devices"
)
import sys
sys.exit(0)
model = models.resnet18().to(device)
inputs = torch.randn(5, 3, 224, 224).to(device)
with profile(activities=activities) as prof:
model(inputs)
prof.export_chrome_trace("trace.json")
你可以在 Chrome 追踪查看器（chrome://tracing）中检查被分析的算子和 CUDA/XPU 内核的序列。
6. 检查堆栈追踪#
分析器可用于分析 Python 和 TorchScript 堆栈追踪。
sort_by_keyword = "self_" + device + "_time_total"
with profile(
activities=activities,
with_stack=True,
experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True),
) as prof:
model(inputs)
# Print aggregated stats
print(prof.key_averages(group_by_stack_n=5).table(sort_by=sort_by_keyword, row_limit=2))
输出可能如下所示（省略了部分列）：
------------------------- -----------------------------------------------------------
Name Source Location
------------------------- -----------------------------------------------------------
aten::thnn_conv2d_forward .../torch/nn/modules/conv.py(439): _conv_forward
.../torch/nn/modules/conv.py(443): forward
.../torch/nn/modules/module.py(1051): _call_impl
.../site-packages/torchvision/models/resnet.py(63): forward
.../torch/nn/modules/module.py(1051): _call_impl
aten::thnn_conv2d_forward .../torch/nn/modules/conv.py(439): _conv_forward
.../torch/nn/modules/conv.py(443): forward
.../torch/nn/modules/module.py(1051): _call_impl
.../site-packages/torchvision/models/resnet.py(59): forward
.../torch/nn/modules/module.py(1051): _call_impl
------------------------- -----------------------------------------------------------
Self CPU time total: 34.016ms
Self CUDA time total: 11.659ms
注意两次卷积以及 torchvision/models/resnet.py 脚本中的两个调用点。
（警告：堆栈追踪会增加额外的分析开销。）
7. 使用分析器分析长时间运行的任务#
PyTorch 分析器提供了一个额外的 API 来处理长时间运行的任务（例如训练循环）。追踪所有的执行过程可能会很慢，并导致生成非常大的追踪文件。为了避免这种情况，请使用可选参数：
schedule- 指定一个函数，该函数接收一个整数参数（步数）作为输入并返回分析器的操作。使用此参数的最佳方法是使用torch.profiler.schedule辅助函数，它可以为你生成调度方案；on_trace_ready- 指定一个函数，该函数接收分析器的引用作为输入，并在每次新的追踪准备就绪时由分析器调用。
为了说明该 API 的工作原理，让我们首先考虑使用 torch.profiler.schedule 辅助函数的以下示例：
分析器假设长时间运行的任务由若干步组成，从零开始编号。上面的示例为分析器定义了以下操作序列：
参数
skip_first告诉分析器忽略前 10 步（skip_first的默认值为零）；在完成前
skip_first步后，分析器开始执行分析循环；每个循环包含三个阶段：
空闲（
wait=5步），在此阶段分析器不处于活动状态；预热（
warmup=1步），在此阶段分析器开始追踪，但结果会被丢弃；此阶段用于丢弃分析器在追踪开始时获取的样本，因为它们通常会受到额外开销的偏差影响；主动追踪（
active=3步），在此阶段分析器会追踪并记录数据；
可选参数
repeat指定了循环次数的上限。默认情况下（值为零），只要任务在运行，分析器就会一直执行循环。
因此，在上面的示例中，分析器将跳过前 15 步，在接下来的 1 步中进行预热，在接下来的 3 步中主动记录，跳过另外 5 步，在接下来的 1 步中预热，再主动记录 3 步。由于指定了 repeat=2 参数，分析器将在前两个循环后停止记录。
在每个循环结束时，分析器会调用指定的 on_trace_ready 函数并将自身作为参数传入。此函数用于处理新的追踪——通过获取表格输出或将输出保存到磁盘作为追踪文件。
要向分析器发送信号以启动下一步，请调用 prof.step() 函数。当前的分析步数存储在 prof.step_num 中。
以下示例展示了如何将上述所有概念用于 CUDA 和 XPU 内核：
sort_by_keyword = "self_" + device + "_time_total"
def trace_handler(p):
output = p.key_averages().table(sort_by=sort_by_keyword, row_limit=10)
print(output)
p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")
with profile(
activities=activities,
schedule=torch.profiler.schedule(wait=1, warmup=1, active=2),
on_trace_ready=trace_handler,
) as p:
for idx in range(8):
model(inputs)
p.step()
了解更多#
查看以下教程以继续学习：