torch.profiler¶
概述¶
PyTorch 分析器是一个工具，允许在训练和推理过程中收集性能指标。 通过分析器的上下文管理器 API，可以更好地了解哪些模型操作最耗费资源， 检查它们的输入形状和堆栈跟踪，研究设备内核活动并可视化执行轨迹。
注意
早期版本的API在torch.autograd模块中被认为是旧版，并将被弃用。
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)[source]¶
低级分析工具包装了自动微分分析器
- Parameters
活动 (可迭代) – 要用于分析的活动组列表（CPU、CUDA），支持的值：
torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA,torch.profiler.ProfilerActivity.XPU。 默认值：ProfilerActivity.CPU 和（如有可用）ProfilerActivity.CUDA 或（如有可用）ProfilerActivity.XPU。record_shapes (bool) – 保存关于算子输入形状的信息。
内存配置文件 (bool) – 跟踪张量内存的分配/释放（详见
export_memory_timeline获取更多细节)。with_stack (bool) – 记录操作的来源信息（文件和行号）。
with_flops (bool) – 使用公式估算特定算子（矩阵乘法和2D卷积）的FLOPS。
with_modules (bool) – 记录模块层次结构（包括函数名称） 对应于操作调用栈。例如，如果模块 A 的前向调用中包含模块 B 的前向调用， 而模块 B 中有一个 aten::add 操作， 那么 aten::add 的模块层次结构为 A.B。 请注意，目前此功能仅支持 TorchScript 模型，而不支持即时模式模型。
实验性配置 (_ExperimentalConfig) – 由像Kineto这样的分析库使用的实验选项集。注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – PyTorch 执行跟踪观察者对象。 PyTorch 执行跟踪 提供基于图的 AI/ML 工作负载表示，并支持重放基准测试、模拟器和仿真器。 当包含此参数时，将在与 PyTorch 分析器相同的时间窗口内调用观察者的 start() 和 stop() 方法。
注意
该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的枚举之一。
输出：以 gzipped JSON、JSON 或 HTML 格式写入的内存时间线。
- 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)[source]¶
性能分析器上下文管理器。
- Parameters
活动 (可迭代) – 要用于分析的活动组列表（CPU、CUDA），支持的值：
torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA,torch.profiler.ProfilerActivity.XPU。 默认值：ProfilerActivity.CPU 和（如有可用）ProfilerActivity.CUDA 或（如有可用）ProfilerActivity.XPU。schedule (Callable) – 可调用对象，接受 step（整数）作为唯一参数，并返回
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 的前向调用中包含模块 B 的前向调用， 而模块 B 中有一个 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 时，分析器会暂时持有张量的引用； 这可能会进一步阻止依赖于引用计数的某些优化，并引入额外的张量复制。
Examples:
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 接口的对象。