注意
点击 此处 下载完整的示例代码
使用 TensorBoard 的 PyTorch Profiler¶
本教程演示如何将 TensorBoard 插件与 PyTorch Profiler 结合使用，以检测模型的性能瓶颈。
简介¶
PyTorch 1.8 包含一个更新的性能分析器 API，它能够记录 CPU 端操作以及 GPU 端的 CUDA 内核启动。性能分析器可以在 TensorBoard 插件中可视化此信息，并提供对性能瓶颈的分析。
在本教程中，我们将使用一个简单的 Resnet 模型来演示如何使用 TensorBoard 插件分析模型性能。
步骤¶
准备数据和模型
使用性能分析器记录执行事件
运行性能分析器
使用 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. 使用性能分析器记录执行事件¶
性能分析器通过上下文管理器启用，并接受多个参数，其中一些最有用的是
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 算子的性能。
“自身”持续时间不包括其子算子的时间。“总计”持续时间包括其子算子的时间。
查看调用栈
单击算子的 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”键将时间线向左和向右移动。您可以多次按下这些键，直到看到可读的表示形式。
如果反向算子的“传入流”字段的值为“对应于反向的前向”，则可以单击文本以获取其启动的前向算子。
在这个例子中，我们可以看到以 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)
然后，让我们在左侧的“运行”下拉列表中选择最近分析的运行。
从上图可以看出，步骤时间减少到大约 76 毫秒，而之前的运行为 132 毫秒，并且 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)。“分配时间”、“释放时间”和“持续时间”列的数据可能在事件发生在时间范围之外时丢失。
在内存统计表中，“大小增量”列将所有分配大小相加并减去所有内存释放大小，即此算子后内存使用量的净增量。“自身大小增量”列类似于“大小增量”，但它不计算子算子的分配。关于 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
在容器内，安装安装轮子包所需的任何依赖项。
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**中所述选择不同的视图。例如，以下是**操作符**视图
在编写本节时，**跟踪**视图无法正常工作，并且不显示任何内容。您可以通过在 Chrome 浏览器中输入 chrome://tracing 来解决此问题。
将
~/profiler_tutorial/log/resnet18目录下的trace.json文件复制到 Windows 系统。
如果文件位于远程位置，您可能需要使用 scp 复制文件。
点击**加载**按钮，从浏览器中
chrome://tracing页面加载跟踪 JSON 文件。
如前所述，您可以移动图形并放大和缩小。您还可以使用键盘在时间线上放大和移动。w 和 s 键以鼠标为中心放大，a 和 d 键将时间线向左和向右移动。您可以多次按下这些键，直到看到可读的表示。
