注意

点击这里下载完整示例代码

PyTorch Profiler 和 TensorBoard¶

创建于: 2021 年 4 月 20 日 | 最后更新于: 2024 年 10 月 31 日 | 最后验证于: 2024 年 11 月 5 日

本教程演示了如何将 TensorBoard 插件与 PyTorch Profiler 结合使用，以检测模型的性能瓶颈。

警告

PyTorch profiler 与 TensorBoard 的集成现已弃用。请改用 Perfetto 或 Chrome tracing 来查看 trace.json 文件。在生成 trace 后，只需将 trace.json 文件拖入 Perfetto UI 或 chrome://tracing 即可可视化你的性能分析结果。

引言¶

PyTorch 1.8 包含一个更新的 profiler API，能够记录 CPU 端操作以及 GPU 端的 CUDA 内核启动。Profiler 可以在 TensorBoard 插件中可视化这些信息，并提供性能瓶颈分析。

在本教程中，我们将使用一个简单的 Resnet 模型来演示如何使用 TensorBoard 插件分析模型性能。

设置¶

安装 torch 和 torchvision 使用以下命令

pip install torch torchvision

步骤¶

准备数据和模型
使用 profiler 记录执行事件
运行 profiler
使用 TensorBoard 查看结果并分析模型性能
借助 profiler 提高性能
使用其他高级功能分析性能
额外实践：在 AMD GPU 上对 PyTorch 进行性能分析

1. 准备数据和模型¶

首先，导入所有必要的库

import torch
import torch.nn
import torch.optim
import torch.profiler
import torch.utils.data
import torchvision.datasets
import torchvision.models
import torchvision.transforms as T

然后准备输入数据。对于本教程，我们使用 CIFAR10 数据集。将其转换为所需格式，并使用 DataLoader 加载每个批次。

transform = T.Compose(
    [T.Resize(224),
     T.ToTensor(),
     T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)

接下来，创建 Resnet 模型、损失函数和优化器对象。要在 GPU 上运行，请将模型和损失函数移动到 GPU 设备。

device = torch.device("cuda:0")
model = torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device)
criterion = torch.nn.CrossEntropyLoss().cuda(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()

定义每个输入数据批次的训练步骤。

def train(data):
    inputs, labels = data[0].to(device=device), data[1].to(device=device)
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

2. 使用 profiler 记录执行事件¶

通过上下文管理器启用 profiler，它接受多个参数，其中一些最有用的是

schedule - 一个可调用对象，接受步数 (int) 作为单个参数，并返回在每一步要执行的 profiler 操作。

在本例中，使用 wait=1, warmup=1, active=3, repeat=1 时，profiler 将跳过第一个步骤/迭代，在第二个步骤开始预热，记录接下来的三个迭代，之后 trace 将可用并调用 on_trace_ready（如果设置）。总共，循环重复一次。TensorBoard 插件中将每个循环称为一个“span”。

在 wait 阶段，profiler 被禁用。在 warmup 阶段，profiler 开始追踪，但结果被丢弃。这样做是为了减少性能分析的开销。性能分析开始时的开销较高，容易给结果带来偏差。在 active 阶段，profiler 工作并记录事件。
on_trace_ready - 在每个循环结束时调用的可调用对象；在本例中，我们使用 torch.profiler.tensorboard_trace_handler 生成用于 TensorBoard 的结果文件。性能分析后，结果文件将保存在 ./log/resnet18 目录下。将此目录指定为 logdir 参数，以便在 TensorBoard 中分析性能。
record_shapes - 是否记录算子输入的形状。
profile_memory - 跟踪张量内存分配/释放。注意，对于 PyTorch 1.10 之前的旧版本，如果性能分析时间过长，请禁用此选项或升级到新版本。
with_stack - 记录算子的源信息（文件和行号）。如果在 VS Code 中启动 TensorBoard（参考），点击堆栈帧将导航到特定的代码行。

with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
        record_shapes=True,
        profile_memory=True,
        with_stack=True
) as prof:
    for step, batch_data in enumerate(train_loader):
        prof.step()  # Need to call this at each step to notify profiler of steps' boundary.
        if step >= 1 + 1 + 3:
            break
        train(batch_data)

或者，也支持以下非上下文管理器方式的启动/停止。

prof = torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
        record_shapes=True,
        with_stack=True)
prof.start()
for step, batch_data in enumerate(train_loader):
    prof.step()
    if step >= 1 + 1 + 3:
        break
    train(batch_data)
prof.stop()

3. 运行 profiler¶

运行以上代码。性能分析结果将保存在 ./log/resnet18 目录下。

4. 使用 TensorBoard 查看结果并分析模型性能¶

注意

TensorBoard 插件支持已弃用，因此其中一些功能可能无法按以前的方式工作。请查看替代方案，HTA。

安装 PyTorch Profiler TensorBoard 插件。

pip install torch_tb_profiler

启动 TensorBoard。

tensorboard --logdir=./log

在 Google Chrome 浏览器或 Microsoft Edge 浏览器中打开 TensorBoard 性能分析 URL（不支持 Safari）。

http://localhost:6006/#pytorch_profiler

您将看到如下图所示的 Profiler 插件页面。

概览

概览显示了模型性能的高级摘要。

“GPU 摘要”面板显示了 GPU 配置、GPU 使用率和 Tensor Cores 使用率。在本例中，GPU 利用率较低。这些指标的详细信息在此。

“步骤时间细分”显示了每个步骤在不同执行类别上花费的时间分布。在本例中，您可以看到 DataLoader 的开销很大。

底部的“性能建议”利用性能分析数据自动突出显示可能的瓶颈，并为您提供可操作的优化建议。

您可以在左侧的“视图”下拉列表中更改查看页面。

算子视图

算子视图显示了在主机或设备上执行的每个 PyTorch 算子的性能。

../_static/img/profiler_operator_view.png

“自身”持续时间不包括其子算子的时间。“总计”持续时间包括其子算子的时间。

查看调用栈

点击算子的 View Callstack，将显示具有相同名称但不同调用栈的算子。然后点击此子表中的 View Callstack，将显示调用栈帧。

如果在 VS Code 中启动 TensorBoard（启动指南），点击调用栈帧将导航到特定的代码行。

内核视图

GPU 内核视图显示了所有内核在 GPU 上花费的时间。

使用 Tensor Cores：此内核是否使用 Tensor Cores。

每个 SM 的平均块数：每个 SM 的块数 = 此内核的块数 / 此 GPU 的 SM 数量。如果此数字小于 1，则表示 GPU 多处理器未充分利用。“每个 SM 的平均块数”是此内核名称所有运行的加权平均值，使用每次运行的持续时间作为权重。

平均估计达到占用率：估计达到占用率定义在此列的工具提示中。对于大多数情况，如内存带宽限制的内核，越高越好。“平均估计达到占用率”是此内核名称所有运行的加权平均值，使用每次运行的持续时间作为权重。

跟踪视图

跟踪视图显示了性能分析的算子和 GPU 内核的时间线。您可以选择它来查看如下详细信息。

您可以使用右侧工具栏移动图形并放大/缩小。键盘也可用于在时间线内缩放和移动。'w' 和 's' 键以鼠标为中心进行放大，而 'a' 和 'd' 键则左右移动时间线。您可以多次按下这些键，直到看到可读的表示形式。

如果反向算子的“输入流”字段的值为“forward correspond to backward”，您可以点击文本以获取其启动的前向算子。

../_static/img/profiler_trace_view_fwd_bwd.png

在本例中，我们可以看到以 enumerate(DataLoader) 为前缀的事件花费了大量时间。在此期间的大部分时间里，GPU 都处于空闲状态。因为此函数在主机端加载和转换数据，在此过程中 GPU 资源被浪费了。

5. 借助 profiler 提高性能¶

在“概览”页面的底部，“性能建议”中的提示指出瓶颈是 DataLoader。PyTorch DataLoader 默认使用单进程。用户可以通过设置参数 num_workers 来启用多进程数据加载。此处有更多详细信息。

在本例中，我们按照“性能建议”将 num_workers 设置如下，将不同的名称（例如 ./log/resnet18_4workers）传递给 tensorboard_trace_handler，然后再次运行。

train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=4)

然后，我们在左侧的“运行”下拉列表中选择最近的性能分析运行结果。

从上面的视图中，我们可以发现步长耗时减少到约 76ms，而之前的运行耗时为 132ms，这主要归功于 DataLoader 时间的减少。

从上面的视图中，我们可以看到 enumerate(DataLoader) 的运行时长减少了，并且 GPU 利用率提高了。

6. 使用其他高级功能分析性能¶

内存视图

要分析内存，必须在 torch.profiler.profile 的参数中将 profile_memory 设置为 True。

您可以通过使用 Azure 上的现有示例来尝试此功能

pip install azure-storage-blob
tensorboard --logdir=https://torchtbprofiler.blob.core.windows.net/torchtbprofiler/demo/memory_demo_1_10

profiler 在性能分析期间记录所有内存分配/释放事件以及分配器的内部状态。内存视图由以下所示的三个组件组成。

这些组件从上到下分别是内存曲线图、内存事件表和内存统计表。

可以在“设备”选择框中选择内存类型。例如，“GPU0”表示下表仅显示每个算子在 GPU 0 上的内存使用情况，不包括 CPU 或其他 GPU。

内存曲线显示了内存消耗的趋势。“已分配”曲线显示实际正在使用的总内存，例如张量。在 PyTorch 中，CUDA 分配器和一些其他分配器采用了缓存机制。“保留”曲线显示分配器保留的总内存。您可以在图表上左键单击并拖动以选择所需范围内的事件。

../_static/img/profiler_memory_curve_selecting.png

选择后，这三个组件将更新为仅显示受限时间范围内的信息，以便您可以获得更多详细信息。通过重复此过程，您可以放大到非常精细的细节。右键单击图表将图表重置到初始状态。

../_static/img/profiler_memory_curve_single.png

在内存事件表中，分配和释放事件被配对为一条条目。“operator”列显示导致分配的直接 ATen 算子。请注意，在 PyTorch 中，ATen 算子通常使用 aten::empty 来分配内存。例如，aten::ones 的实现是先调用 aten::empty，然后调用 aten::fill_。仅将算子名称显示为 aten::empty 帮助不大。在这种特殊情况下，它将显示为 aten::ones (aten::empty)。“分配时间”、“释放时间”和“持续时间”列的数据可能会丢失，如果事件发生在所选时间范围之外。

在内存统计表中，“大小增加”列汇总所有分配大小并减去所有内存释放大小，即此算子执行后的内存使用净增加量。“自身大小增加”列与“大小增加”类似，但不计算子算子的分配。关于 ATen 算子的实现细节，有些算子可能会调用其他算子，因此内存分配可能发生在调用栈的任何级别。也就是说，“自身大小增加”仅计算当前调用栈级别的内存使用增加。最后，“分配大小”列汇总所有分配，不考虑内存释放。

分布式视图

该插件现在支持在使用 NCCL/GLOO 作为后端对 DDP 进行性能分析时的分布式视图。

您可以通过使用 Azure 上的现有示例来尝试此功能

pip install azure-storage-blob
tensorboard --logdir=https://torchtbprofiler.blob.core.windows.net/torchtbprofiler/demo/distributed_bert

../_static/img/profiler_distributed_view.png

“计算/通信概览”显示了计算/通信比率及其重叠程度。从这个视图中，用户可以发现工作节点之间的负载均衡问题。例如，如果某个工作节点的计算 + 重叠时间远大于其他工作节点，则可能存在负载均衡问题，或者此工作节点可能是慢节点 (straggler)。

“同步/通信概览”显示了通信效率。“数据传输时间”是实际数据交换的时间。“同步时间”是等待和与其他工作节点同步的时间。

如果某个工作节点的“同步时间”远短于其他工作节点，此工作节点可能是一个慢节点，它的计算负载可能比其他工作节点更重。

“通信操作统计”汇总了每个工作节点中所有通信算子的详细统计信息。

7. 额外实践：在 AMD GPU 上对 PyTorch 进行性能分析¶

AMD ROCm 平台是一个开源软件堆栈，专为 GPU 计算设计，包含驱动程序、开发工具和 API。我们可以在 AMD GPU 上运行上述步骤。在本节中，我们将在安装 PyTorch 之前使用 Docker 安装 ROCm 基础开发镜像。

为了方便示例，让我们创建一个名为 profiler_tutorial 的目录，并将步骤 1 中的代码另存为 test_cifar10.py 保存在此目录中。

mkdir ~/profiler_tutorial
cd profiler_tutorial
vi test_cifar10.py

撰写本文时，ROCm 平台上的 PyTorch 稳定版 (2.1.1) Linux 版本为 ROCm 5.6。

从 Docker Hub 获取安装了正确用户空间 ROCm 版本的基础 Docker 镜像。

它是 rocm/dev-ubuntu-20.04:5.6。

启动 ROCm 基础 Docker 容器

docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 8G -v ~/profiler_tutorial:/profiler_tutorial rocm/dev-ubuntu-20.04:5.6

在容器内，安装安装 wheels 包所需的任何依赖项。

sudo apt update
sudo apt install libjpeg-dev python3-dev -y
pip3 install wheel setuptools
sudo apt install python-is-python3

安装 wheels 包

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6

安装 torch_tb_profiler，然后运行 Python 文件 test_cifar10.py

pip install torch_tb_profiler
cd /profiler_tutorial
python test_cifar10.py

现在，我们有了在 TensorBoard 中查看所需的所有数据

tensorboard --logdir=./log

选择步骤 4 中描述的不同视图。例如，下面是算子视图

../_static/img/profiler_rocm_tensorboard_operartor_view.png

撰写本节时，跟踪视图无法工作且不显示任何内容。您可以通过在 Chrome 浏览器中输入 chrome://tracing 来解决此问题。

将 ~/profiler_tutorial/log/resnet18 目录下的 trace.json 文件复制到 Windows。

如果文件位于远程位置，您可能需要使用 scp 命令复制文件。

点击加载按钮，从浏览器中的chrome://tracing页面加载追踪JSON文件。

../_static/img/profiler_rocm_chrome_trace_view.png

如前所述，您可以移动图表并进行缩放。您还可以使用键盘在时间轴内进行缩放和移动。w和s键以鼠标为中心进行缩放，而a和d键则左右移动时间轴。您可以多次按下这些键，直到看到清晰可读的表示。

了解更多¶

请查看以下文档以继续学习，并随时在此处提交问题。

脚本总运行时间： ( 0 分钟 0.000 秒)

由 Sphinx-Gallery 生成的图库