torch.profiler¶
概述¶
PyTorch Profiler 是一个工具，允许在训练和推理过程中收集性能指标。Profiler 的上下文管理器 API 可用于更好地了解哪些模型运算符最昂贵，检查其输入形状和调用堆栈，研究设备内核活动以及可视化执行跟踪。
注意
torch.autograd 模块中早期版本的 API 被视为遗留 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, acc_events=False)[source]¶
低级分析器包装 autograd 分析器
- 参数
activities (可迭代对象) – 用于分析的活动组列表（CPU、CUDA），支持的值：
torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA、torch.profiler.ProfilerActivity.XPU。默认值：ProfilerActivity.CPU 以及（如果可用）ProfilerActivity.CUDA 或（如果可用）ProfilerActivity.XPU。record_shapes (bool) – 保存有关运算符输入形状的信息。
profile_memory (bool) – 跟踪张量内存分配/释放（有关更多详细信息，请参见
export_memory_timeline）。with_stack (bool) – 为运算符记录源信息（文件和行号）。
with_flops (bool) – 使用公式估计特定运算符的 FLOPS（矩阵乘法和二维卷积）。
with_modules (bool) – 记录与运算符调用堆栈相对应的模块层次结构（包括函数名称）。例如，如果模块 A 的正向调用调用包含 aten::add 运算符的模块 B 的正向，则 aten::add 的模块层次结构为 A.B。请注意，目前此支持仅适用于 TorchScript 模型，而不适用于渴望模式模型。
experimental_config (_ExperimentalConfig) – Kineto 等分析器库使用的一组实验性选项。请注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – PyTorch Execution Trace Observer 对象。 PyTorch Execution Traces 提供了 AI/ML 工作负载的基于图形的表示，并支持重播基准、模拟器和仿真器。当包含此参数时，观察器的 start() 和 stop() 将与 PyTorch 分析器相同的时间窗口内调用。
acc_events (bool) – 启用跨多个分析周期的 FunctionEvent 的累积
注意
此 API 处于实验阶段，将来可能会发生变化。
启用形状和调用堆栈跟踪会导致额外的开销。当指定 record_shapes=True 时，分析器将暂时保留对张量的引用；这可能会进一步阻止依赖引用计数的某些优化，并引入额外的张量副本。
- export_memory_timeline(path, device=None)[source]¶
从性能分析器收集的树中导出给定设备的内存事件信息，并导出时间线图。使用
export_memory_timeline可以导出 3 个文件，每个文件都由path的后缀控制。对于与 HTML 兼容的绘图，使用后缀
.html，内存时间线图将作为 PNG 文件嵌入到 HTML 文件中。对于由
[times, [sizes by category]]组成的绘图点，其中times是时间戳，sizes是每个类别的内存使用量。内存时间线图将保存为 JSON (.json) 或压缩的 JSON (.json.gz)，具体取决于后缀。对于原始内存点，使用后缀
.raw.json.gz。每个原始内存事件将由(timestamp, action, numbytes, category)组成，其中action是[PREEXISTING, CREATE, INCREMENT_VERSION, DESTROY]中的一个，category是torch.profiler._memory_profiler.Category中的枚举之一。
输出：内存时间线，以压缩的 JSON、JSON 或 HTML 格式写入。
- key_averages(group_by_input_shape=False, group_by_stack_n=0)[source]¶
对事件进行平均，并按运算符名称以及（可选）输入形状和堆栈进行分组。
注意
要使用形状/堆栈功能，请确保在创建性能分析器上下文管理器时设置 record_shapes/with_stack。
- toggle_collection_dynamic(enable, activities)[source]¶
在收集过程中的任何时间点开启/关闭活动的收集。目前支持在 Kineto 中切换 Torch Ops（CPU）和 CUDA 活动。
- 参数
activities (iterable) – 用于性能分析的活动组列表，支持的值：
torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA。
示例
with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ] ) as p: code_to_profile_0() // turn off collection of all CUDA activity p.toggle_collection_dynamic(False, [torch.profiler.ProfilerActivity.CUDA]) code_to_profile_1() // turn on collection of all CUDA activity p.toggle_collection_dynamic(True, [torch.profiler.ProfilerActivity.CUDA]) code_to_profile_2() print(p.key_averages().table( sort_by="self_cuda_time_total", row_limit=-1))
- 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, acc_events=False, use_cuda=None)[source]¶
性能分析器上下文管理器。
- 参数
activities (可迭代对象) – 用于分析的活动组列表（CPU、CUDA），支持的值：
torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA、torch.profiler.ProfilerActivity.XPU。默认值：ProfilerActivity.CPU 以及（如果可用）ProfilerActivity.CUDA 或（如果可用）ProfilerActivity.XPU。schedule (Callable) – 可调用对象，接受 step（int）作为单个参数，并返回
ProfilerAction值，该值指定在每个步骤执行的性能分析器操作。on_trace_ready (Callable) – 在性能分析过程中
schedule返回ProfilerAction.RECORD_AND_SAVE时，在每个步骤调用该可调用对象。record_shapes (bool) – 保存有关运算符输入形状的信息。
profile_memory (bool) – 跟踪张量内存分配/释放。
with_stack (bool) – 为运算符记录源信息（文件和行号）。
with_flops (bool) – 使用公式估算特定运算符（矩阵乘法和二维卷积）的 FLOPs（浮点运算）。
with_modules (bool) – 记录与运算符调用堆栈相对应的模块层次结构（包括函数名称）。例如，如果模块 A 的正向调用调用包含 aten::add 运算符的模块 B 的正向，则 aten::add 的模块层次结构为 A.B。请注意，目前此支持仅适用于 TorchScript 模型，而不适用于渴望模式模型。
experimental_config (_ExperimentalConfig) – 用于 Kineto 库功能的一组实验选项。请注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – 一个 PyTorch 执行跟踪观察器对象。 PyTorch 执行跟踪 提供了 AI/ML 工作负载的基于图的表示形式，并支持重放基准测试、模拟器和仿真器。当包含此参数时，观察器的 start() 和 stop() 将与 PyTorch 性能分析器相同的时间窗口内调用。请参阅下面的示例部分，了解代码示例。
acc_events (bool) – 启用跨多个分析周期的 FunctionEvent 的累积
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函数# Non-default profiler schedule allows user to turn profiler on and off # on different iterations of the training loop; # trace_handler is called every time a new trace becomes available 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, ], # In this example with wait=1, warmup=1, active=2, repeat=1, # profiler will skip the first step/iteration, # start warming up on the second, record # the third and the forth iterations, # after which the trace will become available # and on_trace_ready (when set) is called; # the cycle repeats starting with the next step 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') # used when outputting for tensorboard ) as p: for iter in range(N): code_iteration_to_profile(iter) # send a signal to the profiler that the next iteration has started 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 接口的对象。
