torch.profiler#

创建于: 2020年12月18日 | 最后更新于: 2025年06月13日

概述#

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, acc_events=False, custom_trace_id_callback=None)[source]#

低级分析器包装 autograd profile

参数

activities (iterable) – 要在分析中使用的一组活动（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) – 记录 ops 的源信息（文件和行号）。
with_flops (bool) – 使用公式估算特定运算符（矩阵乘法和二维卷积）的 FLOPs。
with_modules (bool) – 记录与 op 的调用堆栈相对应的模块层次结构（包括函数名称）。例如，如果模块 A 的 forward 调用了包含 aten::add op 的模块 B 的 forward，则 aten::add 的模块层次结构为 A.B。请注意，目前此支持仅适用于 TorchScript 模型，而不适用于 eager 模式模型。
experimental_config (_ExperimentalConfig) – 由 Kineto 等分析器库使用的一组实验选项。注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – PyTorch Execution Trace Observer 对象。PyTorch Execution Traces 提供基于图的 AI/ML 工作负载表示，并支持重放基准测试、模拟器和仿真器。当包含此参数时，将在与 PyTorch profiler 相同的窗口时间内调用 observer 的 start() 和 stop()。
acc_events (bool) – 启用跨多个分析周期的 FunctionEvents 的累积

注意

此 API 处于实验阶段，未来可能会发生更改。

启用形状和堆栈跟踪会产生额外的开销。当指定 record_shapes=True 时，分析器将暂时保留对张量的引用；这可能会进一步阻止某些依赖于引用计数和引入额外张量副本的优化。

add_metadata(key, value)[source]#

将带有字符串键和字符串值的用户定义元数据添加到跟踪文件中

add_metadata_json(key, value)[source]#

将带有字符串键和有效 JSON 值的用户定义元数据添加到跟踪文件中

events()[source]#: 返回未聚合的分析器事件列表，用于跟踪回调或在分析完成后使用

export_chrome_trace(path)[source]#

以 Chrome JSON 格式导出收集到的跟踪。如果启用了 kineto，则只导出计划的最后一个周期。

export_memory_timeline(path, device=None)[source]#

导出分析器收集的树中给定设备的内存事件信息，并导出时间线图。有 3 个可导出的文件使用 export_memory_timeline，每个文件由 path 的后缀控制。

要获得与 HTML 兼容的图，请使用后缀 .html，内存时间线图将作为 PNG 文件嵌入到 HTML 文件中。
对于由 [times, [sizes by category]] 组成的图点，其中 times 是时间戳，sizes 是每个类别的内存使用情况。内存时间线图将保存为 JSON (.json) 或 gzip 压缩的 JSON (.json.gz)，具体取决于后缀。
要获取原始内存点，请使用后缀 .raw.json.gz。每个原始内存事件将由 (timestamp, action, numbytes, category) 组成，其中 action 是 [PREEXISTING, CREATE, INCREMENT_VERSION, DESTROY] 中的一个，而 category 是 torch.profiler._memory_profiler.Category 中的一个枚举。

输出：内存时间线以 gzip 压缩 JSON、JSON 或 HTML 的形式写入。

export_stacks(path, metric='self_cpu_time_total')[source]#

将堆栈跟踪保存到文件

参数

path (str) – 将堆栈文件保存到此位置；
metric (str) – 要使用的指标：“self_cpu_time_total” 或 “self_cuda_time_total”

key_averages(group_by_input_shape=False, group_by_stack_n=0, group_by_overload_name=False)[source]#

平均事件，按运算符名称以及（可选）输入形状、堆栈和重载名称进行分组。

注意

要使用形状/堆栈功能，请确保在创建分析器上下文管理器时设置 record_shapes/with_stack。

preset_metadata_json(key, value)[source]#

在分析器未启动时预设用户定义的元数据，并在稍后添加到跟踪文件中。元数据格式为字符串键和有效 JSON 值

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, custom_trace_id_callback=None)[source]#

分析器上下文管理器。

参数

activities (iterable) – 要在分析中使用的一组活动（CPU、CUDA），支持的值：torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA、torch.profiler.ProfilerActivity.XPU。默认值：ProfilerActivity.CPU 和（如果可用）ProfilerActivity.CUDA 或（如果可用）ProfilerActivity.XPU。
schedule (Callable) – 可调用对象，它接受步数（int）作为单个参数，并返回 ProfilerAction 值，该值指定在每个步数执行的分析器操作。
on_trace_ready (Callable) – 在分析期间，当 schedule 返回 ProfilerAction.RECORD_AND_SAVE 时，在每个步数调用的可调用对象。
record_shapes (bool) – 保存有关运算符输入形状的信息。
profile_memory (bool) – 跟踪张量内存分配/去分配。
with_stack (bool) – 记录 ops 的源信息（文件和行号）。
with_flops (bool) – 使用公式估算特定运算符（矩阵乘法和二维卷积）的 FLOPs（浮点运算）。
with_modules (bool) – 记录与 op 的调用堆栈相对应的模块层次结构（包括函数名称）。例如，如果模块 A 的 forward 调用了包含 aten::add op 的模块 B 的 forward，则 aten::add 的模块层次结构为 A.B。请注意，目前此支持仅适用于 TorchScript 模型，而不适用于 eager 模式模型。
experimental_config (_ExperimentalConfig) – 用于 Kineto 库功能的实验选项集。注意，不保证向后兼容性。
execution_trace_observer (ExecutionTraceObserver) – PyTorch Execution Trace Observer 对象。PyTorch Execution Traces 提供基于图的 AI/ML 工作负载表示，并支持重放基准测试、模拟器和仿真器。当包含此参数时，将在与 PyTorch profiler 相同的窗口时间内调用 observer 的 start() 和 stop()。请参阅下面的示例部分以获取代码示例。
acc_events (bool) – 启用跨多个分析周期的 FunctionEvents 的累积
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 Plugin

注意

示例

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 (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 接口的对象。

get_trace_id()[source]#: 返回当前的跟踪 ID。

set_custom_trace_id_callback(callback)[source]#

设置一个回调函数，在生成新的跟踪 ID 时调用。

step()[source]#

通知分析器下一个分析步骤已开始。

class torch.profiler.ProfilerAction(value)[source]#: 在指定的时间间隔可以执行的分析器操作

class torch.profiler.ProfilerActivity#

成员

CPU

XPU

MTIA

CUDA

HPU

PrivateUse1

property name#

torch.profiler.schedule(*, wait, warmup, active, repeat=0, skip_first=0, skip_first_wait=0)[source]#

返回一个可调用对象，可用作分析器的 schedule 参数。分析器将跳过前 skip_first 步，然后等待 wait 步，然后进行 warmup 步的预热，然后进行 active 步的活动记录，然后以 wait 步开始重复循环。可选的循环次数由 repeat 参数指定，零值表示循环将一直持续到分析完成。

skip_first_wait 参数控制是否应跳过第一个 wait 阶段。这在用户希望在 skip_first 之后等待比 skip_first 更长的时间，但又不希望在第一个配置文件中等待时很有用。例如，如果 skip_first 为 10，wait 为 20，如果 skip_first_wait 为零，第一个循环将在预热前等待 10 + 20 = 30 步，但如果 skip_first_wait 非零，则只等待 10 步。之后的所有循环将在最后一个活动和预热之间等待 20 步。

返回类型: Callable

torch.profiler.tensorboard_trace_handler(dir_name, worker_name=None, use_gzip=False)[source]#

将跟踪文件输出到 dir_name 目录，然后该目录可以直接作为 logdir 提供给 tensorboard。worker_name 在分布式场景中应对于每个 worker 都是唯一的，它将默认为 ‘[hostname]_[pid]’。

Intel Instrumentation and Tracing Technology APIs#

torch.profiler.itt.is_available()[source]#: 检查 ITT 功能是否可用

torch.profiler.itt.mark(msg)[source]#

描述在某个时间点发生的瞬时事件。

参数: msg (str) – 与事件关联的 ASCII 消息。

torch.profiler.itt.range_push(msg)[source]#

将一个范围推送到嵌套范围跨度的堆栈上。返回已开始范围的零基深度。

参数: msg (str) – 与范围关联的 ASCII 消息

torch.profiler.itt.range_pop()[source]#: 将一个范围从嵌套范围跨度的堆栈上弹出。返回已结束范围的零基深度。

torch.profiler#

概述#

API 参考#

Intel Instrumentation and Tracing Technology APIs#

文档

教程

资源