评价此页
★ ★ ★ ★ ★
intermediate/tensorboard_profiler_tutorial

注意

跳转到结尾 以下载完整的示例代码。

PyTorch Profiler 与 TensorBoard#

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

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

警告

PyTorch profiler 与 TensorBoard 的集成现已弃用。相反，请使用 Perfetto 或 Chrome trace 来查看 trace.json 文件。在 生成 trace 后，只需将 trace.json 拖到 Perfetto UI 或 chrome://tracing 中以可视化您的 profile。

简介#

PyTorch 1.8 包含更新的 profiler API，能够记录 CPU 侧的操作以及 GPU 侧的 CUDA kernel 启动。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 加载每个 batch。

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()

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

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 - 一个可调用对象，它接受 step (int) 作为单个参数，并在每个 step 返回要执行的 profiler 操作。

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

    在 wait step 期间，profiler 被禁用。在 warmup step 期间，profiler 开始 tracing 但结果会被丢弃。这是为了减少 profiling 开销。Profiling 开始时的开销很高，很容易使 profiling 结果产生偏差。在 active step 期间，profiler 工作并记录事件。

  • on_trace_ready - 一个在每个循环结束时调用的可调用对象；在本例中，我们使用 torch.profiler.tensorboard_trace_handler 生成 TensorBoard 的结果文件。Profiling 之后，结果文件将保存到 ./log/resnet18 目录中。将此目录指定为 logdir 参数，以便在 TensorBoard 中分析 profile。

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

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

  • with_stack - 记录操作的文件和行号等源代码信息。如果 TensorBoard 在 VS Code 中启动 (参考)，则单击堆栈帧将导航到特定代码行。

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)

或者，以下非上下文管理器 start/stop 也受支持。

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#

运行上述代码。Profiling 结果将保存在 ./log/resnet18 目录中。

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

注意

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

安装 PyTorch Profiler TensorBoard 插件。

pip install torch_tb_profiler

启动 TensorBoard。

tensorboard --logdir=./log

在 Google Chrome 浏览器或 Microsoft Edge 浏览器中打开 TensorBoard profile URL (Safari 不受支持)。

http://localhost:6006/#pytorch_profiler

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

  • 概述

../_static/img/profiler_overview1.png

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

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

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

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

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

  • Operator 视图

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

../_static/img/profiler_operator_view.png

“Self” 持续时间不包括其子 operator 的时间。“Total” 持续时间包括其子 operator 的时间。

  • 查看调用堆栈

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

../_static/img/profiler_callstack.png

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

../_static/img/profiler_vscode.png

  • Kernel 视图

GPU kernel 视图显示 GPU 上花费的所有 kernel 的时间。

../_static/img/profiler_kernel_view.png

Tensor Cores Used：此 kernel 是否使用 Tensor Cores。

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

Mean Est. Achieved Occupancy：Est. Achieved Occupancy 在此列的工具提示中定义。在大多数情况下，例如受内存带宽限制的 kernel，越高越好。“Mean Est. Achieved Occupancy” 是所有此 kernel 名称的运行的加权平均值，使用每个运行的持续时间作为权重。

  • Trace 视图

Trace 视图显示 profiled operator 和 GPU kernel 的时间线。您可以选择它以查看详细信息，如下所示。

../_static/img/profiler_trace_view1.png

您可以使用右侧工具栏移动图形并放大/缩小。还可以使用键盘在时间线内进行缩放和移动。多次按这些键，直到看到可读的表示形式。

如果 backward operator 的“Incoming Flow”字段的值为“forward correspond to backward”，您可以单击该文本以获取其启动的 forward operator。

../_static/img/profiler_trace_view_fwd_bwd.png

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

5. 借助 profiler 改进性能#

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

在本例中，我们遵循“Performance Recommendation”，并将 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”下拉列表中选择最近 profiled 的运行。

../_static/img/profiler_overview2.png

从上面的视图中，我们可以发现 step 时间减少到大约 76 毫秒，与之前运行的 132 毫秒相比，并且 DataLoader 的时间减少主要贡献。

../_static/img/profiler_trace_view2.png

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

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

  • Memory 视图

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

您可以尝试在 Azure 上使用现有示例

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

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

../_static/img/profiler_memory_view.png

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

可以在“Device”选择框中选择内存类型。例如，“GPU0”表示下表仅显示每个 operator 在 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

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

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

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

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

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

在撰写本文时，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 来解决此问题。

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

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

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

../_static/img/profiler_rocm_chrome_trace_view.png

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

了解更多#

请查看以下文档以继续学习，并随时在此处打开 issue。

在此页面上
PyTorch 库