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, acc_events=False, custom_trace_id_callback=None)

低级分析工具包装了自动微分分析器

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() 方法。

• acc_events (bool) – 启用在多个分析周期中累积FunctionEvents

注意

该API处于试验阶段，未来可能会发生变化。

启用形状和堆栈跟踪会导致额外的开销。 当指定 record_shapes=True 时，分析器会暂时持有张量的引用； 这可能会进一步阻止依赖于引用计数的某些优化，并引入额外的张量复制。

add_metadata(key, value)

将一个由用户定义的元数据（包含字符串键和字符串值）添加到跟踪文件中

add_metadata_json(key, value)

添加一个用户定义的元数据，其中键为字符串，值为有效的 JSON 数据到跟踪文件中

events()

返回未聚合的分析事件列表， 用于跟踪回调或在分析完成后使用。

export_chrome_trace(path)

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

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。每个原始内存 事件将由 (timestamp, action, numbytes, category) 组成，其中 action 是 [PREEXISTING, CREATE, INCREMENT_VERSION, DESTROY] 中的一个， 而 category 是来自 torch.profiler._memory_profiler.Category 的枚举之一。

输出：以 gzipped JSON、JSON 或 HTML 格式写入的内存时间线。

export_stacks(path, metric='self_cpu_time_total')

将堆栈跟踪保存到文件

Parameters

• path (str) – 将堆栈文件保存到此位置;

• metric (str) – 要使用的度量: “self_cpu_time_total” 或 “self_cuda_time_total”

key_averages(group_by_input_shape=False, group_by_stack_n=0)

按操作符名称和（可选）输入形状及堆栈对事件进行分组并求平均值。

注意

要使用 shape/stack 功能，请确保在创建性能分析器上下文管理器时设置 record_shapes/with_stack。

preset_metadata_json(key, value)

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

toggle_collection_dynamic(enable, activities)

在收集过程中的任何时间点，都可以开启或关闭活动的收集。目前支持切换 Torch Ops（CPU）和 Kineto 中支持的 CUDA 活动

Parameters

activities (iterable) – 用于性能分析的活动组列表，支持的值: torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA

Examples:

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)

性能分析器上下文管理器。

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 分析器相同的时间窗口内被调用。请参阅下面的示例部分以获取代码示例。

• 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 插件

注意

启用形状和堆栈跟踪会导致额外的开销。 当指定 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 接口的对象。

get_trace_id()

返回当前的跟踪ID。

set_custom_trace_id_callback(callback)

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

step()

表示分析器下一个分析步骤已开始。

class torch.profiler.ProfilerAction(value)

可在指定间隔执行的分析器操作

class torch.profiler.ProfilerActivity

Members:

CPU

XPU

MTIA

CUDA

PrivateUse1

property name

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

返回一个可调用对象，可用作分析器 schedule 参数。分析器将跳过前 skip_first 步，然后等待 wait 步，接着对接下来的 warmup 步进行预热， 然后对接下来的 active 步进行活动记录，之后从 wait 步开始重复该周期。可以通过 repeat 参数指定可选的周期数，零值表示周期将持续到分析完成为止。

The skip_first_wait 参数控制是否跳过第一个 wait 阶段。 如果用户希望在循环之间等待的时间长于 skip_first 步，但不适用于第一个配置文件，则此功能可能很有用。例如，如果 skip_first 是 10 而 wait 是 20，在 skip_first_wait 为零的情况下，第一个循环将在开始预热前等待 10 + 20 = 30 步；而在 skip_first_wait 为非零的情况下，只会等待 10 步。所有后续循环则会在最后一个活跃阶段和预热之间等待 20 步。

Return type

可调用的

torch.profiler.tensorboard_trace_handler(dir_name, worker_name=None, use_gzip=False)

将跟踪文件输出到dir_name的目录中，然后该目录可以直接作为logdir传递给tensorboard。 在分布式场景中，worker_name对每个工作进程应该是唯一的，默认情况下会设置为‘[hostname]_[pid]’。

英特尔仪器与跟踪技术API

torch.profiler.itt.is_available()

检查 ITT 功能是否可用

torch.profiler.itt.mark(msg)

描述某个时刻发生的一个瞬间事件。

Parameters

msg (str) – 与事件相关联的ASCII消息。

torch.profiler.itt.range_push(msg)

将一个范围推送到嵌套范围跨度的堆栈上。返回所启动范围的零基深度。

Parameters

msg (str) – 要与范围关联的ASCII消息

torch.profiler.itt.range_pop()

从嵌套范围跨度的堆栈中弹出一个范围。返回结束的范围的以零为基数的深度。