22、PYTORCH PROFILER
这个方法解释了如何使用PyTorch分析器并测量模型操作符的时间和内存消耗。
Introduction
PyTorch包含一个简单的分析器API，当用户需要确定模型中最昂贵的操作符时，这个API非常有用。
在本教程中，我们将使用一个简单的Resnet模型来演示如何使用分析器来分析模型性能。
Setup
使用以下命令安装torch和torchvision:
pip install torch torchvision
Steps
1、导入所有必要的库
2、实例化一个简单的Resnet模型
3、使用分析器分析执行时间
4、使用分析器分析内存消耗
5、使用跟踪功能
6、检查堆栈跟踪
7、将数据可视化为火焰图
8、使用分析器分析长时间运行的作业
1. Import all necessary libraries
在这个方法中，我们将使用torch, torchvision.models和profiler modules:
import torch import torchvision.models as models from torch.profiler import profile, record_function, ProfilerActivity
2. Instantiate a simple Resnet model
让我们创建一个Resnet模型的实例并为它准备一个输入:
model = models.resnet18() inputs = torch.randn(5, 3, 224, 224)
3. Using profiler to analyze execution time
PyTorch分析器是通过上下文管理器启用的，它接受许多参数，其中最有用的有:
activities - 要分析的活动列表：
ProfilerActivity.CPU - PyTorch 操作符、TorchScript 函数和用户定义的代码标签（参见下面的 record_function）；
ProfilerActivity.CUDA - 设备上的 CUDA 内核；
record_shapes - 是否记录操作员输入的形状；
profile_memory - 是否报告模型张量消耗的内存量；
use_cuda - 是否测量 CUDA 内核的执行时间。
注意：使用 CUDA 时，分析器还会显示主机上发生的运行时 CUDA 事件。
让我们看看我们如何使用分析器来分析执行时间：
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 启动的异步任务和（在向后传递的情况下）使用向后（）调用启动的向后传递运算符。
让我们打印出上面执行的统计信息：
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.549ms
在这里我们看到，正如预期的那样，大部分时间都花在了卷积上（特别是在使用 MKL-DNN 支持编译的 PyTorch 的 mkldnn_convolution 中）。 注意self cpu time和cpu time的区别——运营商可以调用其他运营商，self cpu time不包括子运营商调用花费的时间，而总cpu time包括它。 您可以通过将 sort_by="self_cpu_time_total" 传递到 table 调用中来选择按 self 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))
# (omitting some columns) # --------------------------------- ------------ ------------------------------------------- # 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 上执行的模型的性能：
model = models.resnet18().cuda()
inputs = torch.randn(5, 3, 224, 224).cuda()
with profile(activities=[
ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
with record_function("model_inference"):
model(inputs)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
（注意：第一次使用 CUDA 分析可能会带来额外的开销。）
结果表输出：
# (omitting some columns) # ------------------------------------------------------- ------------ ------------ # 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
请注意输出中出现的设备上内核（例如 sgemm_32x32x32 NN）。
4. Using profiler to analyze memory consumption
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))
# (omitting some columns)
# --------------------------------- ------------ ------------ ------------
# 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. Using tracing functionality
分析结果可以作为 .json 跟踪文件输出：
model = models.resnet18().cuda()
inputs = torch.randn(5, 3, 224, 224).cuda()
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
model(inputs)
prof.export_chrome_trace("trace.json")
您可以在 Chrome 跟踪查看器 (chrome://tracing) 中检查分析运算符和 CUDA 内核的序列：
6. Examining stack traces
Profiler 可用于分析 Python 和 TorchScript 堆栈跟踪：
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
with_stack=True,
) as prof:
model(inputs)
# Print aggregated stats
print(prof.key_averages(group_by_stack_n=5).table(sort_by="self_cuda_time_total", row_limit=2))
# (omitting some columns)
# ------------------------- -----------------------------------------------------------
# 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. Visualizing data as a flamegraph
执行时间（self_cpu_time_total 和 self_cuda_time_total 指标）和堆栈跟踪也可以可视化为火焰图。 为此，首先使用 export_stacks 导出原始数据（需要 with_stack=True）：
prof.export_stacks("/tmp/profiler_stacks.txt", "self_cuda_time_total")
我们建议使用例如 生成交互式 SVG 的 Flamegraph 工具：
# git clone https://github.com/brendangregg/FlameGraph # cd FlameGraph # ./flamegraph.pl --title "CUDA time" --countname "us." /tmp/profiler_stacks.txt > perf_viz.svg
8. Using profiler to analyze long-running jobs
PyTorch 分析器提供了一个额外的 API 来处理长时间运行的作业（例如训练循环）。 跟踪所有执行可能会很慢并导致非常大的跟踪文件。 为避免这种情况，请使用可选参数：
schedule - 指定一个函数，它接受一个整数参数（步骤号）作为输入并返回一个分析器的动作，使用这个参数的最好方法是使用 torch.profiler.schedule 辅助函数，它可以为你生成一个时间表；
on_trace_ready - 指定一个函数，它将对探查器的引用作为输入，并在每次新跟踪准备好时由探查器调用。
为了说明 API 的工作原理，让我们首先考虑以下带有 torch.profiler.schedule 辅助函数的示例：
from torch.profiler import schedule
my_schedule = schedule(
skip_first=10,
wait=5,
warmup=1,
active=3,
repeat=2)
Profiler 假定长时间运行的作业由步骤组成，从零开始编号。 上面的示例为探查器定义了以下操作序列：
1、参数skip_first 告诉profiler 它应该忽略前10 个步骤（skip_first 的默认值为零）；
2、在第一个 skip_first 步骤之后，探查器开始执行探查器周期；
3、每个周期由三个阶段组成：
空闲（等待 = 5 步），在此阶段分析器未激活；
预热（warmup=1 步），在此阶段分析器开始跟踪，但结果被丢弃； 此阶段用于丢弃在跟踪开始时分析器获得的样本，因为它们通常会因额外的开销而倾斜；
主动跟踪（主动=3 步），在此阶段分析器跟踪和记录数据；
4、可选的重复参数指定循环数的上限。 默认情况下（零值），只要作业运行，探查器就会执行循环。
因此，在上面的示例中，profiler 将跳过前 15 个步骤，将下一步花在热身上，主动记录接下来的 3 个步骤，再跳过 5 个步骤，将下一步花在热身上，再主动记录另外 3 个步骤 . 由于指定了 repeat=2 参数值，分析器将在前两个循环后停止记录。
在每个周期结束时，分析器调用指定的 on_trace_ready 函数并将其自身作为参数传递。 此函数用于处理新的跟踪 - 通过获取表输出或将输出保存在磁盘上作为跟踪文件。
要向分析器发送下一步已开始的信号，请调用 prof.step() 函数。 当前分析器步骤存储在 prof.step_num 中。
以下示例显示了如何使用上述所有概念：
def trace_handler(p):
output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)
print(output)
p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
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()