PyTorch Profiler 与 TensorBoard 结合使用

1. 准备数据和模型

首先，导入所有必要的库：

importtorch
importtorch.nn
importtorch.optim
importtorch.profiler
importtorch.utils.data
importtorchvision.datasets
importtorchvision.models
importtorchvision.transformsasT

然后准备输入数据。在本教程中，我们使用 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()

为每批输入数据定义训练步骤。

deftrain(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 插件中称为一个“span”。

  在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 - 记录操作的源信息（文件和行号）。如果 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)

此外，还支持以下非上下文管理器的启动/停止操作。

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 Core 使用率。在此示例中，GPU 利用率较低。这些指标的详细信息请参见 此处。

“步骤时间分解”展示了每个步骤在不同执行类别中的时间分布情况。在此示例中，可以看到 DataLoader 的开销较大。

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

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

• 操作员视图

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

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

• 查看调用栈

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

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

• 内核视图

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

Tensor Core 使用情况：该内核是否使用了 Tensor Core。

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

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

• 跟踪视图

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

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

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

在此示例中，我们可以看到前缀为 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”下拉列表中选择最近分析的运行。

从上述视图中，我们可以看到步骤时间从之前的 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

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

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

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

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

选择后，三个组件将更新为限制的时间范围，以便您可以获取更多相关信息。通过重复此过程，您可以深入到非常细粒度的细节。右键点击图表将重置图表到初始状态。

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

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

• 分布式视图

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

您可以通过在 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）为 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

• 安装 wheel 包：

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视图：

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

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

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

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

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