Table of Contents
  • Tutorials >
  • 使用TensorBoard的PyTorch性能分析工具
Shortcuts
intermediate/tensorboard_profiler_tutorial

备注

点击:ref:`此处 <sphx_glr_download_intermediate_tensorboard_profiler_tutorial.py>`下载完整示例代码

使用TensorBoard的PyTorch性能分析工具

Created On: Apr 20, 2021 | Last Updated: Oct 31, 2024 | Last Verified: Nov 05, 2024

此教程演示了如何使用TensorBoard插件与PyTorch性能分析工具来检测模型的性能瓶颈。

警告

与PyTorch性能分析工具集成的TensorBoard现已弃用。相反，使用Perfetto或Chrome追踪查看``trace.json``文件。在`生成追踪文件`（<https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html#using-tracing-functionality>）之后，只需将``trace.json``拖到`Perfetto UI`（<https://ui.perfetto.dev/>）或``chrome://tracing``中即可可视化你的配置文件。

介绍

PyTorch 1.8引入了更新的性能分析工具API，它可以记录CPU端操作以及GPU端的CUDA内核启动。性能分析工具可以在TensorBoard插件中可视化这些信息并提供性能瓶颈分析。

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

设置

安装``torch``和``torchvision``可以使用以下命令：

pip install torch torchvision

步骤

  1. 准备数据和模型

  2. 使用性能分析工具记录执行事件

  3. 运行性能分析工具

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

  5. 借助性能分析工具提升性能

  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. 使用性能分析工具记录执行事件

性能分析工具通过上下文管理器启用，并接受多个参数，其中一些最有用的是：

  • schedule - 一个可调用对象，它以步长（int）作为单个参数，并返回每一步执行的性能分析操作。

    在此示例中，设置 wait=1, warmup=1, active=3, repeat=1，分析器将跳过第一个步骤/迭代，从第二个开始预热，记录接下来的三个迭代，之后分析结果将可用，并调用设置的 on_trace_ready（如果设置了）。总体而言，此循环重复一次。在 TensorBoard 插件中，每个循环称为一个“跨度”。

    在 wait 步骤期间，分析器被禁用。在 warmup 步骤期间，分析器开始跟踪但结果会被丢弃，这是为了减少分析开销。在分析开始时的开销较高，容易对分析结果产生偏差。在 active 步骤期间，分析器正常工作并记录事件。

  • on_trace_ready - 每个循环结束时调用的函数；在此示例中，我们使用 torch.profiler.tensorboard_trace_handler 来为 TensorBoard 生成结果文件。分析完成后，结果文件将保存到 ./log/resnet18 目录。指定此目录作为 logdir 参数以在 TensorBoard 中分析配置文件。

  • record_shapes - 是否记录操作输入的形状。

  • profile_memory - 跟踪张量内存分配/释放。注意，对于版本低于 1.10 的旧版本 PyTorch，如果分析时间过长，请禁用它或升级到新版本。

  • 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 分析器 TensorBoard 插件。

pip install torch_tb_profiler

启动 TensorBoard。

tensorboard --logdir=./log

在 Google Chrome 浏览器或 Microsoft Edge 浏览器中打开 TensorBoard 配置文件 URL（不支持 Safari）。

http://localhost:6006/#pytorch_profiler

您可以看到 Profiler 插件页面，如下所示。

  • 概述

../_static/img/profiler_overview1.png

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

“GPU 概况”面板显示 GPU 配置、GPU 使用率和 Tensor Cores 使用情况。在本示例中，GPU 使用率较低。这些指标的详细信息在`此处 <https://github.com/pytorch/kineto/blob/main/tb_plugin/docs/gpu_utilization.md>`_。

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

底部的“性能建议”使用分析数据自动突出可能的瓶颈，并提供可执行的优化建议。

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

  • 操作符视图

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

../_static/img/profiler_operator_view.png

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

  • 查看调用堆栈

单击操作符的 View Callstack，将显示具有相同名称但不同调用堆栈的操作符。然后在该子表中单击 View Callstack，调用堆栈帧将被显示。

../_static/img/profiler_callstack.png

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

../_static/img/profiler_vscode.png

  • 内核视图

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

../_static/img/profiler_kernel_view.png

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

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

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

  • 跟踪视图

跟踪视图显示被分析的操作符和 GPU 内核的时间轴。您可以选择它以查看更多详情，如下所示。

../_static/img/profiler_trace_view1.png

通过右侧工具栏可以移动图表并放大/缩小，同时键盘也可用于缩放和在时间轴中移动。“w”和“s”键围绕鼠标中心放大，“a”和“d”键左右移动时间轴。您可以多次按这些键，直到看到易于读取的表示。

如果某个后向操作符的“传入流量”字段的值为“前向对应后向”，您可以单击文本以获得其启动的前向操作符。

../_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)

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

../_static/img/profiler_overview2.png

从上述视图中，我们可以发现步骤时间减少到约 76 毫秒，与之前运行的 132 毫秒相比下降了很多，而时间减少主要归功于 DataLoader。

../_static/img/profiler_trace_view2.png

从上述视图中，我们可以看到 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

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

../_static/img/profiler_memory_view.png

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

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

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

../_static/img/profiler_memory_curve_selecting.png

选择后，三个组件将针对所选时间范围更新，以便您获得更多信息。重复此过程，可缩放至非常详细的水平。右键单击图表将重置图表到初始状态。

../_static/img/profiler_memory_curve_single.png

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

“计算/通信概览”显示计算/通信比例及其重叠程度。从该视图中，用户可以发现工人的负载平衡问题。例如，如果一个工人的计算+重叠时间远大于其他工人，则可能存在负载平衡问题或该工人可能是慢速工人。

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

如果一个工人的“同步时间”远短于其他工人的，此工人可能是一个慢速工人，可能计算负载高于其他工人。

“通信操作统计”总结每个工人中所有通信操作的详细统计信息。

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

  • 安装 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 中所描述。例如，下面是 Operator 视图：

../_static/img/profiler_rocm_tensorboard_operartor_view.png

在撰写本节内容时，Trace 视图无法正常工作，显示为空。可以通过在 Chrome 浏览器中输入 chrome://tracing 来解决。

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

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

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

../_static/img/profiler_rocm_chrome_trace_view.png

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

文档

访问 PyTorch 的详细开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

为了分析流量并优化您的体验，我们在此网站上提供Cookie。点击或浏览即表示您同意我们使用Cookie。 作为本网站的当前维护者，Facebook的Cookie政策适用。了解更多信息，包括可用的控件： Cookies政策.