注意
跳转至页面底部下载完整示例代码。
PyTorch Profiler#
创建日期：2021年1月29日 | 最后更新：2025年7月9日 | 最后验证：未验证
本指南介绍了如何使用 PyTorch Profiler 以及如何测量模型算子的时间和内存消耗。
简介#
PyTorch 包含一个简单的 Profiler API，当用户需要确定模型中最耗时的算子时，它非常有用。
在本指南中，我们将使用一个简单的 Resnet 模型来演示如何使用 Profiler 分析模型性能。
先决条件#
torch >= 2.3.0
设置#
要安装 torch 和 torchvision，请使用以下命令
pip install torch torchvision
步骤#
导入所有必要的库
实例化一个简单的 Resnet 模型
使用 Profiler 分析执行时间
使用 Profiler 分析内存消耗
使用追踪（Tracing）功能
检查堆栈跟踪
使用 Profiler 分析长期运行的作业
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. 使用 Profiler 分析执行时间#
PyTorch Profiler 通过上下文管理器启用，并接受多个参数，其中一些最常用的参数包括：
activities- 需要分析的活动列表ProfilerActivity.CPU- PyTorch 算子、TorchScript 函数和用户定义的代码标签（见下文record_function）；ProfilerActivity.CUDA- 设备上的 CUDA 内核；ProfilerActivity.XPU- 设备上的 XPU 内核；
record_shapes- 是否记录算子输入张量的形状；profile_memory- 是否报告模型张量消耗的内存量；
注意：使用 CUDA 时，Profiler 还会显示主机上发生的 CUDA 运行时事件。
让我们看看如何使用 Profiler 来分析执行时间
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
with record_function("model_inference"):
model(inputs)
请注意，我们可以使用 record_function 上下文管理器，用用户提供的名称标记任意代码段（在上面的示例中使用了 model_inference 作为标签）。
Profiler 允许用户检查在被其上下文管理器包裹的代码段执行期间调用了哪些算子。如果同时激活了多个 Profiler 范围（例如在并行的 PyTorch 线程中），每个分析上下文管理器仅追踪其对应范围内的算子。Profiler 还会自动追踪使用 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（注意：这要求在运行 Profiler 时设置 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 出现了两次，且输入形状不同。
Profiler 也可以用于分析在 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. 使用 Profiler 分析内存消耗#
PyTorch Profiler 还可以显示模型算子执行期间分配（或释放）的内存量（模型张量所使用的内存）。在下面的输出中，“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. 检查堆栈跟踪#
Profiler 可用于分析 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. 使用 Profiler 分析长期运行的作业#
PyTorch Profiler 提供了额外的 API 来处理长期运行的作业（例如训练循环）。追踪整个执行过程可能会很慢，并导致产生非常大的追踪文件。为避免这种情况，请使用可选参数
schedule- 指定一个函数，该函数接收一个整数参数（步数）作为输入并为分析器返回一个动作。使用此参数的最佳方式是使用torch.profiler.schedule辅助函数，它可以为你生成调度方案；on_trace_ready- 指定一个函数，该函数接收 Profiler 的引用作为输入，并在每次新追踪就绪时被调用。
为了说明 API 的工作原理，我们首先考虑以下使用 torch.profiler.schedule 辅助函数的示例
from torch.profiler import schedule
my_schedule = schedule(skip_first=10, wait=5, warmup=1, active=3, repeat=2)
Profiler 假设长期运行的作业由一系列从 0 开始编号的步骤组成。上述示例为 Profiler 定义了以下动作序列：
skip_first参数告诉 Profiler 应该忽略前 10 个步骤（skip_first的默认值为 0）；在完成前
skip_first个步骤后，Profiler 开始执行周期；每个周期包含三个阶段：
空闲（
wait=5步），在此阶段 Profiler 不处于活动状态；预热（
warmup=1步），在此阶段 Profiler 开始追踪，但结果会被丢弃；此阶段用于丢弃分析器在追踪开始时获取的样本，因为它们通常会被额外开销所扭曲；主动追踪（
active=3步），在此阶段 Profiler 追踪并记录数据；
可选的
repeat参数指定了周期的上限。默认情况下（值为 0），Profiler 将在作业运行期间持续执行周期。
因此，在上面的示例中，Profiler 将跳过前 15 步，用接下来的 1 步进行预热，积极记录接下来的 3 步，再跳过 5 步，再用 1 步进行预热，积极记录另外 3 步。由于指定了 repeat=2 参数值，Profiler 将在前两个周期后停止记录。
在每个周期结束时，Profiler 会调用指定的 on_trace_ready 函数并将自身作为参数传递。该函数用于处理新的追踪——通过获取表格输出或将输出作为追踪文件保存到磁盘。
要向 Profiler 发送下一阶段已开始的信号，请调用 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()
了解更多#
请查阅以下指南/教程以继续学习