性能分析器
性能分析器是一个工具，允许在训练和推理过程中收集性能指标。性能分析器的上下文管理器 API 可用于更好地了解哪些模型运算符最昂贵，检查它们的输入形状和堆栈跟踪，研究设备内核活动，并可视化执行跟踪。它提供了对模型性能的洞察，使您能够对其进行优化和改进。
本指南介绍如何使用 PyTorch Profiler 来测量模型算子的时间和内存消耗，以及如何将其与 Accelerate 集成。我们将涵盖各种用例并为每个用例提供示例。
使用 Profiler 分析执行时间
Profiler 允许检查在用 Profiler 上下文管理器包装的代码范围内执行期间调用的算子。
让我们看看如何使用 Profiler 分析执行时间
PyTorch
加速
import torch
import torchvision.models as models
from torch.profiler import profile, record_function, ProfilerActivity
model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
model(inputs)
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))结果表格输出（省略了一些列）
--------------------------------- ------------ ------------ ------------ ------------
Name Self CPU CPU total CPU time avg # of Calls
--------------------------------- ------------ ------------ ------------ ------------
aten::conv2d 171.000us 52.260ms 2.613ms 20
aten::convolution 227.000us 52.089ms 2.604ms 20
aten::_convolution 270.000us 51.862ms 2.593ms 20
aten::mkldnn_convolution 51.273ms 51.592ms 2.580ms 20
aten::batch_norm 118.000us 7.059ms 352.950us 20
aten::_batch_norm_impl_index 315.000us 6.941ms 347.050us 20
aten::native_batch_norm 6.305ms 6.599ms 329.950us 20
aten::max_pool2d 40.000us 4.008ms 4.008ms 1
aten::max_pool2d_with_indices 3.968ms 3.968ms 3.968ms 1
aten::add_ 780.000us 780.000us 27.857us 28
--------------------------------- ------------ ------------ ------------ ------------
Self CPU time total: 67.016ms为了获得更细粒度的结果并包含算子的输入形状，请传递 group_by_input_shape=True（注意：这需要使用 record_shapes=True 运行 Profiler）
print(prof.key_averages(group_by_input_shape=True).table(sort_by="cpu_time_total", row_limit=10))使用 Profiler 分析内存消耗
Profiler 还可以显示在模型算子执行期间分配（或释放）的内存量（模型张量使用的内存）。要启用内存分析功能，请传递 profile_memory=True。
PyTorch
加速
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))结果表格输出（省略了一些列）
--------------------------------- ------------ ------------ ------------
Name CPU Mem Self CPU Mem # of Calls
--------------------------------- ------------ ------------ ------------
aten::empty 94.85 Mb 94.85 Mb 205
aten::max_pool2d_with_indices 11.48 Mb 11.48 Mb 1
aten::addmm 19.53 Kb 19.53 Kb 1
aten::mean 10.00 Kb 10.00 Kb 1
aten::empty_strided 492 b 492 b 5
aten::cat 240 b 240 b 6
aten::abs 480 b 240 b 4
aten::masked_select 120 b 112 b 1
aten::ne 61 b 53 b 3
aten::eq 30 b 30 b 1
--------------------------------- ------------ ------------ ------------
Self CPU time total: 69.332ms导出 Chrome 追踪
您可以在 Chrome 追踪查看器（chrome://tracing）中检查已分析算子和 CUDA 内核的顺序。
PyTorch
加速
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")使用 Profiler 分析长时间运行的任务
Profiler 提供了一个额外的 API 来处理长时间运行的任务（例如训练循环）。跟踪所有执行可能会很慢并导致非常大的追踪文件。为了避免这种情况，请使用可选参数
schedule_option：调度选项允许您控制何时启用分析。这对于长时间运行的任务很有用，可以避免收集过多的数据。可用的键为wait、warmup、active、repeat和skip_first。Profiler 将跳过前skip_first步，然后等待wait步，然后进行接下来的warmup步的预热，然后进行接下来的active步的主动记录，然后重复从wait步开始的循环。可选的循环次数由repeat参数指定，零值表示循环将持续到分析完成。on_trace_ready：指定一个函数，该函数将 Profiler 的引用作为输入，并在每次新的追踪准备好时由 Profiler 调用。
为了说明 API 的工作原理，请考虑以下示例
PyTorch
加速
from torch.profiler import schedule
my_schedule = schedule(
skip_first=10,
wait=5,
warmup=1,
active=3,
repeat=2
)
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=my_schedule,
on_trace_ready=trace_handler
) as p:
for idx in range(8):
model(inputs)
p.step()FLOPS
使用公式估算特定算子的 FLOPS（浮点运算）（矩阵乘法和二维卷积）。
测量浮点运算 (FLOPS)
PyTorch
加速
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
with_flops=True
) as prof:
model(inputs)
print(prof.key_averages().table(sort_by="flops", row_limit=10))结果表格输出（省略了一些列）
------------------------------------------------------- ------------ ------------ ------------
Name Self CPU Self CUDA Total FLOPs
------------------------------------------------------- ------------ ------------ ------------
aten::conv2d 197.000us 0.000us 18135613440.000
aten::addmm 103.000us 17.000us 5120000.000
aten::mul 29.000us 2.000us 30.000
aten::convolution 409.000us 0.000us --
aten::_convolution 253.000us 0.000us --
aten::cudnn_convolution 5.465ms 2.970ms --
cudaEventRecord 138.000us 0.000us --
cudaStreamIsCapturing 43.000us 0.000us --
cudaStreamGetPriority 40.000us 0.000us --
cudaDeviceGetStreamPriorityRange 10.000us 0.000us --
------------------------------------------------------- ------------ ------------ ------------
Self CPU time total: 21.938ms
Self CUDA time total: 4.165ms总结和更多信息
PyTorch Profiler 是一个用于分析模型性能的强大工具。通过将其与 Accelerate 集成，您可以轻松地分析您的模型并深入了解其性能，从而帮助您优化和改进模型。
有关更详细的信息，请参阅 PyTorch Profiler 文档。
< > 在 GitHub 上更新