torch.profiler¶
概述¶
PyTorch Profiler 是一个工具，允许在训练和推理期间收集性能指标。 Profiler 的上下文管理器 API 可以用来更好地理解哪些模型操作是最昂贵的， 检查它们的输入形状和堆栈跟踪，研究设备内核活动并可视化执行跟踪。
注意
在torch.autograd模块中的API的早期版本被视为遗留版本，并将被弃用。
API参考¶
- class torch.profiler._KinetoProfile(*, activities=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None, execution_trace_observer=None)[源代码]¶
低级分析器包装了自动求导分析
- Parameters
活动（可迭代）——用于分析的活动组（CPU、CUDA）列表，支持的值：
torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA。 默认值：ProfilerActivity.CPU 和（当可用时）ProfilerActivity.CUDA。record_shapes (bool) – 保存关于操作符输入形状的信息。
profile_memory (bool) – 跟踪张量的内存分配/释放（更多详情请参见
export_memory_timeline）。with_stack (bool) – 记录操作的源信息（文件和行号）。
with_flops (bool) – 使用公式来估计特定操作符的FLOPS（矩阵乘法和2D卷积）。
with_modules (bool) – 记录与操作的调用栈相对应的模块层次结构（包括函数名称）。例如，如果模块 A 的 forward 调用了模块 B 的 forward，其中包含一个 aten::add 操作，那么 aten::add 的模块层次结构是 A.B。请注意，目前此支持仅存在于 TorchScript 模型中，而不适用于急切模式模型。
experimental_config (_ExperimentalConfig) – 一组实验性选项，用于像Kineto这样的分析器库。注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – 一个PyTorch执行跟踪观察器对象。 PyTorch执行跟踪提供了一种基于图形的AI/ML工作负载表示，并支持回放基准测试、模拟器和仿真器。 当包含此参数时，观察器的start()和stop()将在与PyTorch分析器相同的时间窗口内被调用。
注意
此API是实验性的，未来可能会发生变化。
启用形状和堆栈跟踪会导致额外的开销。 当指定 record_shapes=True 时，分析器将暂时保留对张量的引用； 这可能会进一步阻止依赖于引用计数的某些优化，并引入额外的张量复制。
- export_memory_timeline(path, device=None)[源代码]¶
从分析器收集的树中导出给定设备的内存事件信息，并导出时间线图。使用
export_memory_timeline可以导出 3 个文件，每个文件由path的后缀控制。对于兼容HTML的图表，使用后缀
.html，内存时间线图将以PNG文件的形式嵌入到HTML文件中。对于由
[times, [sizes by category]]组成的绘图点，其中times是时间戳，sizes是每个类别的内存使用情况。 内存时间线图将保存为 JSON (.json) 或压缩的 JSON (.json.gz)，具体取决于后缀。对于原始内存点，使用后缀
.raw.json.gz。每个原始内存事件将包含(时间戳, 动作, 字节数, 类别)，其中动作是[预先存在的, 创建, 增加版本, 销毁]之一， 并且类别是来自torch.profiler._memory_profiler.Category的枚举之一。
输出：内存时间线以gzip压缩的JSON、JSON或HTML格式写入。
- export_stacks(path, metric='self_cpu_time_total')[源代码]¶
将堆栈跟踪保存到适合可视化的文件格式中。
- Parameters
注意
使用 FlameGraph 工具的示例：
cd FlameGraph
./flamegraph.pl –title “CPU 时间” –countname “微秒” profiler.stacks > perf_viz.svg
- class torch.profiler.profile(*, activities=None, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None, execution_trace_observer=None, use_cuda=None)[源代码]¶
分析器上下文管理器。
- Parameters
活动（可迭代）——用于分析的活动组（CPU、CUDA）列表，支持的值：
torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA。 默认值：ProfilerActivity.CPU 和（当可用时）ProfilerActivity.CUDA。schedule (可调用对象) – 一个可调用对象，接受步数（int）作为单一参数并返回指定在每一步执行的分析器动作的
ProfilerAction值。on_trace_ready (可调用对象) – 在每次
schedule返回ProfilerAction.RECORD_AND_SAVE时调用的可调用对象。record_shapes (bool) – 保存关于操作符输入形状的信息。
profile_memory (bool) – 跟踪张量内存的分配/释放。
with_stack (bool) – 记录操作的源信息（文件和行号）。
with_flops (bool) – 使用公式来估计特定操作符（矩阵乘法和2D卷积）的浮点运算次数（FLOPs）。
with_modules (bool) – 记录与操作的调用栈相对应的模块层次结构（包括函数名称）。例如，如果模块 A 的 forward 调用了模块 B 的 forward，其中包含一个 aten::add 操作，那么 aten::add 的模块层次结构是 A.B。请注意，目前此支持仅存在于 TorchScript 模型中，而不适用于急切模式模型。
experimental_config (_ExperimentalConfig) – 一组用于Kineto库特性的实验选项。注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – 一个 PyTorch 执行跟踪观察器对象。 PyTorch 执行跟踪提供了一个基于图形的 AI/ML 工作负载表示，并支持回放基准测试、模拟器和仿真器。 当包含此参数时，观察器的 start() 和 stop() 将在与 PyTorch 分析器相同的时间窗口内被调用。请参阅下面的示例部分以获取代码示例。
use_cuda (bool) –
自版本1.8.1起已弃用: 请改用
activities。
注意
使用
schedule()来生成可调用的调度。 非默认调度在分析长时间训练任务时非常有用，并且允许用户在训练过程的不同迭代中获取多个跟踪。 默认调度只是连续记录上下文管理器持续时间内的所有事件。注意
使用
tensorboard_trace_handler()生成用于 TensorBoard 的结果文件：on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name)分析完成后，结果文件可以在指定的目录中找到。使用以下命令：
tensorboard --logdir dir_name查看 TensorBoard 中的结果。 更多信息，请参阅 PyTorch Profiler TensorBoard 插件
注意
启用形状和堆栈跟踪会导致额外的开销。 当指定 record_shapes=True 时，分析器将暂时保留对张量的引用； 这可能会进一步阻止依赖于引用计数的某些优化，并引入额外的张量复制。
示例：
with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ] ) as p: code_to_profile() print(p.key_averages().table( sort_by="self_cuda_time_total", row_limit=-1))
使用分析器的
schedule、on_trace_ready和step函数：```python# 非默认的分析器调度允许用户在训练循环的不同迭代中打开和关闭分析器； # trace_handler 每次有新的跟踪可用时都会被调用 def trace_handler(prof): print(prof.key_averages().table( sort_by="self_cuda_time_total", row_limit=-1)) # prof.export_chrome_trace("/tmp/test_trace_" + str(prof.step_num) + ".json") with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], # 在这个例子中，wait=1, warmup=1, active=2, repeat=1， # 分析器将跳过第一个步骤/迭代， # 在第二个步骤开始预热，记录 # 第三个和第四个迭代， # 之后跟踪将变为可用 # 并且当设置时，on_trace_ready 被调用； # 循环从下一个步骤开始重复 schedule=torch.profiler.schedule( wait=1, warmup=1, active=2, repeat=1), on_trace_ready=trace_handler # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log') # 用于输出到 tensorboard 时使用 ) as p: for iter in range(N): code_iteration_to_profile(iter) # 向分析器发送信号，表示下一个迭代已经开始 p.step()
```以下示例展示了如何设置执行跟踪观察器（execution_trace_observer）
with torch.profiler.profile( ... execution_trace_observer=( ExecutionTraceObserver().register_callback("./execution_trace.json") ), ) as p: for iter in range(N): code_iteration_to_profile(iter) p.step()
您也可以参考 tests/profiler/test_profiler.py 中的 test_execution_trace_with_kineto()。 注意：也可以传递任何满足 _ITraceObserver 接口的对象。
