注意

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

PyTorch Profiler With TensorBoard¶

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

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

警告

PyTorch 分析器与 TensorBoard 的集成现已弃用。请改用 Perfetto 或 Chrome 跟踪来查看 trace.json 文件。在生成跟踪文件后，只需将 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 记录执行事件
运行分析器
使用 TensorBoard 查看结果并分析模型性能
借助分析器提升性能
利用其他高级功能分析性能
附加实践：在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记录执行事件¶

分析器通过上下文管理器启用，并接受几个参数，其中一些最有用的包括：

schedule - 可调用对象，接受 step (int) 作为单个参数并返回在每个步骤中要执行的分析器操作。

在这个示例中，wait=1, warmup=1, active=3, repeat=1， profiler 会跳过第一个步骤/迭代，在第二个步骤开始预热，记录接下来的三个迭代，之后追踪数据将变为可用，并调用 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. 运行分析器¶

运行上面的代码。性能分析结果将保存在 ./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 的开销相当显著。

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

你可以在左侧“Views”下拉列表中更改视图页面。

操作符视图

操作员视图显示了在主机或设备上执行的每个 PyTorch 操作符的性能。

../_static/img/profiler_operator_view.png

“Self” 时间不包括其子操作符的时间。 “Total” 时间包括其子操作符的时间。

查看调用栈

点击一个操作符的 View Callstack，将显示同名但不同调用栈的操作符。然后点击此子表中的 View Callstack，将显示调用栈帧。

如果在VS Code内部启动TensorBoard (启动指南)，单击调用堆栈帧将导航到特定代码行。

内核视图

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

使用的Tensor Core：此内核是否使用Tensor Core。

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

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

跟踪视图

跟踪视图显示了已分析操作符和 GPU 内核的时间线。你可以选择它以查看以下详细信息。

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

如果一个反向操作符的“Incoming Flow”字段的值为“forward correspond to backward”，你可以点击该文本以获取其启动的正向操作符。

../_static/img/profiler_trace_view_fwd_bwd.png

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

5. 通过分析器提高性能¶

在“概述”页面底部，“性能建议”中的提示表明瓶颈是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)

那么，让我们选择左侧“Runs”下拉列表中最近配置的运行。

从上面的视图中，我们可以发现步骤时间减少到约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

分析器会记录所有内存分配/释放事件以及分配器的内部状态。内存视图由以下三个部分组成，如下所示。

组件包括内存曲线图、内存事件表和内存统计表，从上到下依次排列。

内存类型可以在“设备”选择框中进行选择。例如，“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)。如果事件发生在时间范围之外，“Allocation Time”、“Release Time”和“Duration”列的数据可能会缺失。

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

分布式视图

该插件现在支持在使用 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

“计算/通信概览”展示了计算与通信的比例及其重叠程度。从这个视角，用户可以识别出各个工作节点之间的负载平衡问题。例如，如果某个工作节点的计算时间加上重叠时间明显大于其他节点，则可能存在负载不平衡的问题，或者该工作节点可能是一个拖后腿的节点。

“同步/通信概览”展示了通信的效率。 “数据传输时间”是实际数据交换所花费的时间。 “同步时间”是等待并与其他工作者同步所花费的时间。

如果某个工作者的“同步时间”明显短于其他工作者的，该工作者可能是拖后腿的节点，其计算工作量可能比其他工作者更多。

“通信操作统计”汇总了每个工作进程中所有通信操作的详细统计信息。

7. 额外实践：在AMD GPU上配置PyTorch¶

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

为了举例说明，让我们创建一个名为profiler_tutorial的目录，并将步骤1中的代码保存到此目录下的test_cifar10.py中。

mkdir ~/profiler_tutorial
cd profiler_tutorial
vi test_cifar10.py

在撰写本文时，PyTorch 在 ROCm 平台上的稳定(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

安装轮子：

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

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

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

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

点击 Load 按钮从浏览器的 chrome://tracing 页面加载 trace JSON 文件。

../_static/img/profiler_rocm_chrome_trace_view.png

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

了解更多¶

查看以下文档以继续您的学习，并随时在这里提出问题。

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

通过 Sphinx-Gallery 生成的画廊