6.3 PyTorch Profiler (Profiler)

文档摘要

6.3 PyTorch Profiler (Profiler) 6.3 PyTorch Profiler (Profiler) 详解与实践在深度学习模型的开发和优化过程中，性能分析至关重要。了解模型的瓶颈所在，可以帮助我们有效地提升训练和推理速度，降低资源消耗。PyTorch Profiler (以下简称 Profiler) 是 PyTorch 生态系统中一个强大的性能分析工具，它能够帮助开发者深入了解 PyTorch 程序的性能特征，从而进行针对性的优化。本篇文章将深入探讨 PyTorch Profiler 的各个方面，包括其核心概念、配置选项、使用方法、代码实践以及结果分析。通过本文，您将能够掌握如何有效地利用 Profiler 来提升 PyTorch 模型的性能。 6.3.

6.3 PyTorch Profiler (Profiler)

6.3 PyTorch Profiler (Profiler) 详解与实践

在深度学习模型的开发和优化过程中，性能分析至关重要。了解模型的瓶颈所在，可以帮助我们有效地提升训练和推理速度，降低资源消耗。PyTorch Profiler (以下简称 Profiler) 是 PyTorch 生态系统中一个强大的性能分析工具，它能够帮助开发者深入了解 PyTorch 程序的性能特征，从而进行针对性的优化。

本篇文章将深入探讨 PyTorch Profiler 的各个方面，包括其核心概念、配置选项、使用方法、代码实践以及结果分析。通过本文，您将能够掌握如何有效地利用 Profiler 来提升 PyTorch 模型的性能。

6.3.1 Profiler 的重要性与基本概念

在深度学习模型训练过程中，我们常常会遇到以下问题：

训练速度慢： 模型训练耗时过长，影响开发效率和迭代速度。
资源利用率低： GPU 或 CPU 等硬件资源没有得到充分利用，造成浪费。
内存溢出： 模型或数据超出内存容量，导致程序崩溃。
性能瓶颈不明： 不清楚模型的哪个部分是性能瓶颈，难以进行优化。

Profiler 的出现正是为了解决这些问题。它通过监控和记录 PyTorch 程序运行时的各种事件和指标，例如：

操作符 (Operators) 的执行时间： 哪些操作符耗时最多？
GPU 和 CPU 的利用率： 硬件资源是否饱和？
内存分配和释放： 是否存在内存泄漏或不必要的内存分配？
函数调用栈： 代码的执行路径是怎样的？
CUDA kernel 的执行情况： GPU kernel 的耗时和效率如何？

通过分析 Profiler 收集到的数据，我们可以清晰地了解程序的性能瓶颈，并采取相应的优化措施，例如：

优化模型结构： 减少模型参数量，简化网络结构。
优化数据加载： 提升数据加载速度，避免数据瓶颈。
优化 CUDA kernel： 针对性地优化耗时较长的 CUDA kernel。
调整超参数： 优化学习率、batch size 等超参数。
使用更高效的操作符或算法： 替换低效的操作符或算法。

简而言之，Profiler 是深度学习模型性能优化的必备工具。它可以帮助开发者从“盲人摸象”式的优化转变为基于数据驱动的精准优化。

6.3.2 PyTorch Profiler 的核心组件与工作流程

PyTorch Profiler 主要由以下几个核心组件构成：

torch.profiler.profile (Context Manager): 这是 Profiler 的核心入口，通过上下文管理器的方式来启动和停止性能分析。我们需要将需要分析的代码块放在 with torch.profiler.profile(...) 语句块中。
torch.profiler.record_function (Decorator/Context Manager): 用于更细粒度地分析代码中的特定函数或代码块。可以使用装饰器或上下文管理器的方式来标记需要记录的函数或代码块。
schedule (Scheduling Configuration): 控制 Profiler 的记录行为，例如设置预热 (warmup) 阶段、活动 (active) 阶段、等待 (wait) 阶段和跳过 (skip) 阶段，以灵活地控制 Profiler 的启动和停止时机。
on_trace_ready (Callback Function): Profiler 完成数据收集后，会调用 on_trace_ready 指定的回调函数来处理收集到的性能数据。我们可以自定义回调函数来保存 trace 文件、生成报告或者进行其他分析操作。
profile_memory (Boolean): 控制是否记录内存分配和释放事件。开启后可以分析内存使用情况。
record_shapes (Boolean): 控制是否记录操作符的输入张量的形状信息。开启后可以了解操作符处理的数据形状。
with_stack (Boolean): 控制是否记录函数调用栈信息。开启后可以追踪代码的执行路径。
activities (List of torch.profiler.ProfilerActivity): 指定需要记录的活动类型，例如 CPU 活动、CUDA 活动等。

Profiler 的工作流程可以用以下 Mermaid 图表示：

流程说明：

程序启动后，当代码执行到 with torch.profiler.profile(...) 语句块时，Profiler 初始化。
Profiler 根据 schedule 配置 (如果存在) 控制记录状态，否则持续记录。
在 Profiler 记录期间，会收集 CPU、GPU、内存、操作符等性能数据。
当 with torch.profiler.profile(...) 语句块结束时，Profiler 停止记录。
Profiler 根据 on_trace_ready 配置 (如果存在) 调用回调函数处理性能数据，否则进行默认处理。
回调函数 (或默认处理) 负责保存 trace 文件、生成报告等操作。
程序继续运行，性能分析完成。

6.3.3 Profiler 的基本使用：代码实践

接下来，我们通过一些代码示例来演示 PyTorch Profiler 的基本使用方法。

示例 1：使用 profile 上下文管理器进行基本性能分析


import torch
import torch.nn as nn
import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity
# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(128, 10)
    def forward(self, x):
        x = x.view(x.size(0), -1) # flatten
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x
model = SimpleModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 随机生成输入数据和目标
inputs = torch.randn(64, 1, 28, 28)
targets = torch.randint(0, 10, (64,))
# 使用 profile 上下文管理器进行性能分析
with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/simple_model')) as prof:
    with record_function("model_train"):
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        prof.step() # 手动调用 step() 来推进 Profiler 的迭代步数
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
# 也可以将 trace 保存为 Chrome Trace 格式
# prof.export_chrome_trace("./trace.json")

代码解释：

导入必要的库: 导入 torch.profiler 模块以及其他 PyTorch 库。
定义模型和数据: 定义一个简单的线性模型 SimpleModel，以及损失函数 CrossEntropyLoss 和优化器 SGD。随机生成输入数据 inputs 和目标 targets。
使用 profile 上下文管理器:
- activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]: 指定要记录 CPU 和 CUDA 活动。
- record_shapes=True: 记录操作符的输入张量形状。
- on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/simple_model'): 配置 on_trace_ready 回调函数，使用 tensorboard_trace_handler 将 trace 数据写入 TensorBoard 日志文件，日志文件将保存在 ./log/simple_model 目录下。
使用 record_function 上下文管理器:
- with record_function("model_train"):: 使用 record_function 上下文管理器标记模型训练的代码块，方便在 Profiler 结果中区分不同的代码段。
模型训练代码: 包含前向传播、计算损失、反向传播和优化器更新等训练步骤。
prof.step(): 在每次迭代结束后手动调用 prof.step() 来推进 Profiler 的迭代步数，这对于正确记录 schedule 和迭代信息非常重要。
prof.key_averages().table(...): 打印 Profiler 结果的摘要表格，按照 cpu_time_total 排序，显示前 10 行。
prof.export_chrome_trace("./trace.json"): （注释掉的代码）可以将 trace 数据导出为 Chrome Trace 格式的 JSON 文件，可以使用 Chrome 浏览器 chrome://tracing 功能打开进行可视化分析。

运行代码后，您可以在控制台看到类似以下的摘要表格：


------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                              Name    Self CPU %      CPU total %     CPU total         CPU self        CUDA total        CUDA self       # of Calls      Self CPU     CPU total
------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                    aten::relu_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             1      0.000ms      0.001ms
                                           aten::empty_strided_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             1      0.000ms      0.001ms
                                              aten::linear_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             2      0.000ms      0.001ms
                                             aten::view_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             1      0.000ms      0.001ms
                                 aten::empty_like_strided_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             1      0.000ms      0.001ms
                                         aten::addmm_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             2      0.000ms      0.001ms
                                        aten::mm_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             2      0.000ms      0.001ms
                                        aten::t_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             2      0.000ms      0.001ms
                                     aten::as_strided_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             2      0.000ms      0.001ms
                                   aten::size_                               0.0%         13.8%      0.000ms       0.000ms      0.000ms       0.000ms             1      0.000ms      0.001ms
------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------

表格列解释：

Name: 操作符的名称。
Self CPU %: 该操作符自身 CPU 时间占总 CPU 时间的百分比。
CPU total %: 该操作符及其子调用 CPU 时间占总 CPU 时间的百分比。
CPU total: 该操作符及其子调用的总 CPU 时间 (毫秒)。
CPU self: 该操作符自身的 CPU 时间 (毫秒)。
CUDA total: 该操作符及其子调用的总 CUDA 时间 (毫秒)。
CUDA self: 该操作符自身的 CUDA 时间 (毫秒)。
# of Calls: 该操作符被调用的次数。
Self CPU: 该操作符自身的 CPU 时间 (微秒)。
CPU total: 该操作符及其子调用的总 CPU 时间 (微秒)。

TensorBoard 可视化：

运行代码后，您可以使用 TensorBoard 查看 Profiler 生成的日志。在终端中运行以下命令：


tensorboard --logdir ./log/simple_model

然后在浏览器中打开 TensorBoard (通常是 http://localhost:6006)，您可以看到 Profiler 收集的性能数据，包括操作符的时间消耗、内存使用情况、CUDA kernel 执行情况等，以更直观的方式进行分析。

示例 2：使用 schedule 配置 Profiler 行为

在实际训练过程中，我们通常不需要从头到尾都进行性能分析。例如，模型在初始阶段可能处于不稳定的状态，此时的 Profiler 数据可能不具有代表性。我们可以使用 schedule 参数来配置 Profiler 的行为，例如设置预热阶段，只在训练的特定阶段进行详细分析。


import torch
import torch.nn as nn
import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity, schedule, ProfilerAction
# ... (模型和数据定义与示例 1 相同) ...
# 配置 schedule
prof_schedule = schedule(
    wait=1, # 等待 1 步
    warmup=1, # 预热 1 步
    active=3, # 记录 3 步
    repeat=2 # 重复整个 schedule 2 次
)
# 使用 profile 上下文管理器并配置 schedule
with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        schedule=prof_schedule,
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/scheduled_model')) as prof:
    for step in range(10): # 训练 10 步
        with record_function("model_train_step"):
            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        prof.step() # 手动调用 step() 来推进 Profiler 的迭代步数
        print(f"Step {step+1} done")
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

代码解释：

定义 prof_schedule: 使用 schedule 函数定义 Profiler 的记录计划。
- wait=1: Profiler 在开始记录前等待 1 步。
- warmup=1: Profiler 进行 1 步的预热，预热阶段的数据不会被记录到最终的 trace 文件中，但用于 Profiler 内部的预热。
- active=3: Profiler 记录 3 步的性能数据。
- repeat=2: 整个 schedule (wait -> warmup -> active) 重复执行 2 次。
在 profile 中使用 schedule=prof_schedule: 将定义的 prof_schedule 传递给 profile 上下文管理器。

运行结果分析：

在这个示例中，Profiler 将会记录训练过程中的第 3-5 步和第 8-10 步的性能数据。在 TensorBoard 中，您可以看到只有这两个阶段的数据被记录下来。schedule 参数可以帮助我们更精确地控制 Profiler 的记录时机，避免记录不必要的数据，并更专注于分析关键阶段的性能。

示例 3：使用 record_function 进行细粒度分析

record_function 可以用于标记代码中的特定函数或代码块，从而进行更细粒度的性能分析。例如，我们可以分别标记前向传播、反向传播和数据加载等关键部分，以便更清晰地了解各个部分的耗时情况。


import torch
import torch.nn as nn
import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity
# ... (模型和数据定义与示例 1 相同) ...
# 使用 profile 上下文管理器
with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/granular_model')) as prof:
    with record_function("data_loading"): # 标记数据加载
        inputs = torch.randn(64, 1, 28, 28).cuda() # 模拟数据加载耗时
        targets = torch.randint(0, 10, (64,)).cuda()
    with record_function("forward_pass"): # 标记前向传播
        outputs = model(inputs)
    with record_function("loss_calculation"): # 标记损失计算
        loss = criterion(outputs, targets)
    with record_function("backward_pass"): # 标记反向传播
        optimizer.zero_grad()
        loss.backward()
    with record_function("optimizer_step"): # 标记优化器更新
        optimizer.step()
        prof.step()
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

代码解释：

使用多个 record_function 上下文管理器: 分别使用 record_function 上下文管理器标记数据加载、前向传播、损失计算、反向传播和优化器更新等代码块，并赋予不同的名称，例如 "data_loading", "forward_pass" 等。

运行结果分析：

在 TensorBoard 或 Chrome Trace 中查看 Profiler 结果时，您可以看到以 "data_loading", "forward_pass" 等名称标记的事件，可以清晰地了解各个代码段的耗时情况，从而更精细地定位性能瓶颈。

6.3.4 Profiler 的高级特性与最佳实践

除了基本用法外，PyTorch Profiler 还提供了一些高级特性和最佳实践，可以帮助我们更有效地进行性能分析：

内存分析 (profile_memory=True): 开启 profile_memory=True 参数可以记录内存分配和释放事件。通过分析内存数据，可以发现内存泄漏、不必要的内存分配等问题，并进行优化，例如减少中间变量的生成、使用 in-place 操作等。
记录形状信息 (record_shapes=True): 开启 record_shapes=True 参数可以记录操作符的输入张量形状。这对于分析模型中数据形状的变化、理解操作符的行为非常有帮助。
记录函数调用栈 (with_stack=True): 开启 with_stack=True 参数可以记录函数调用栈信息。这可以帮助我们追踪代码的执行路径，尤其是在复杂模型或代码结构中，更容易理解程序的执行流程。
自定义 on_trace_ready 回调函数: on_trace_ready 参数允许我们自定义回调函数来处理 Profiler 收集到的数据。我们可以根据自己的需求，实现更灵活的数据处理和分析逻辑，例如：
- 将 trace 数据保存为自定义格式。
- 生成更详细的性能报告 (例如，统计每个操作符的平均耗时、标准差等)。
- 将 Profiler 数据上传到监控系统。
- 结合其他性能分析工具进行更深入的分析。
使用 NVIDIA Nsight Systems 进行更深入的 GPU 分析: PyTorch Profiler 可以与 NVIDIA Nsight Systems 工具集成，进行更深入的 GPU kernel 级别的性能分析。Nsight Systems 可以提供更详细的 CUDA kernel 执行信息，例如 kernel 的占用率、指令级别的时间消耗等，帮助我们更精细地优化 GPU 性能。
在分布式训练中使用 Profiler: PyTorch Profiler 也支持在分布式训练环境中使用。在分布式训练中，我们需要在每个进程中都启动 Profiler，并将各个进程的 trace 数据收集起来进行分析。可以使用 torch.distributed.get_rank() 获取当前进程的 rank，并在 on_trace_ready 回调函数中根据 rank 来保存不同的 trace 文件。

最佳实践建议：

逐步分析： 先进行粗粒度的分析，例如先分析整个训练过程的性能，找到瓶颈所在的大模块，再逐步细化到模块内部的函数或操作符。
关注关键指标： 在 Profiler 结果中，重点关注 CPU 时间、GPU 时间、内存使用率、操作符耗时等关键指标。
结合领域知识： 结合深度学习领域的知识和模型特点，分析 Profiler 结果，例如，对于卷积神经网络，可以重点关注卷积层和池化层的性能。
迭代优化： 根据 Profiler 的分析结果，进行针对性的优化，然后再次使用 Profiler 进行验证，迭代优化，直到达到满意的性能。

6.3.5 总结

PyTorch Profiler 是一个功能强大的性能分析工具，可以帮助开发者深入了解 PyTorch 程序的性能瓶颈，从而进行精准优化。本文详细介绍了 Profiler 的核心概念、配置选项、使用方法、代码实践以及高级特性和最佳实践。

通过学习本文，您应该已经掌握了如何使用 PyTorch Profiler 来分析和优化您的深度学习模型。在实际开发过程中，建议您积极使用 Profiler，并结合实际情况和领域知识，不断提升模型的性能，提高开发效率。

希望这篇文章能够帮助您更好地理解和使用 PyTorch Profiler！

作者与出处

原作者：灏天文库

来源：灏天文库

原文： https://aiknowledge.cn/collection/215-pytorch%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0%E6%A1%86%E6%9E%B6%E5%AE%9E%E6%88%98

整理：灏天文库整理

由灏天文库平台收录，内容或由平台用户上传，仅供学习交流