intermediate/tensorboard_profiler_tutorial

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

带 TensorBoard 的 PyTorch Profiler

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

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

警告

PyTorch profiler 与 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

步骤

1. 准备数据和模型

2. 使用 profiler 记录执行事件

3. 运行 profiler

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

5. 借助 profiler 提升性能

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

7. 额外实践：在 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 将跳过第一个步骤/迭代，在第二个步骤开始预热，记录接下来的三个迭代，之后跟踪将可用，并调用 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 操作符的性能。

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

• 查看调用堆栈

点击操作符的 查看调用堆栈，将显示同名但调用堆栈不同的操作符。然后点击此子表中的 查看调用堆栈，将显示调用堆栈帧。

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

• 内核视图

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

使用的张量核心数：此内核是否使用张量核心。

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

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

• 跟踪视图

跟踪视图显示了分析过的操作符和 GPU 内核的时间轴。您可以选择它来查看详细信息，如下所示。

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

如果反向运算符的“Incoming Flow”字段值为“forward correspond to backward”，您可以点击文本以获取其启动的正向运算符。

在此示例中，我们可以看到以 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)

然后我们选择左侧“运行”下拉列表中最近分析的运行。

从上图可以看出，与之前运行的 132ms 相比，步长耗时减少到约 76ms，主要贡献是 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 分配器和其他一些分配器采用了缓存机制。“已预留”曲线显示分配器预留的总内存。您可以在图表上左键单击并拖动以选择所需范围内的事件

选择后，三个组件将根据受限的时间范围进行更新，以便您可以获得更多相关信息。通过重复此过程，您可以深入到非常细致的细节。右键单击图形将图形重置为初始状态。

在内存事件表中，分配和释放事件配对成一个条目。“操作符”列显示导致分配的直接 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

“计算/通信概览”显示计算/通信比率及其重叠程度。从此视图中，用户可以找出工作器之间的负载平衡问题。例如，如果一个工作器的计算 + 重叠时间远大于其他工作器，则可能存在负载平衡问题或该工作器可能是滞后者。

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

如果一个工作器的“同步时间”比其他工作器短得多，则该工作器可能是滞后者，其计算工作负载可能比其他工作器多。

“通信操作统计”汇总了每个工作器中所有通信操作的详细统计数据。

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**中描述的不同视图。例如，下面是**操作符**视图

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

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

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

• 单击“加载”按钮以从浏览器中的 chrome://tracing 页面加载跟踪 JSON 文件。

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

了解更多

请查阅以下文档以继续学习，如有任何问题，请随时在此处提交。

• PyTorch TensorBoard Profiler Github

• torch.profiler API

• HTA