注意
点击这里下载完整的示例代码
PyTorch Profiler 与 TensorBoard¶
创建于:2021 年 4 月 20 日 | 最后更新:2024 年 10 月 31 日 | 最后验证:2024 年 11 月 05 日
本教程演示如何将 TensorBoard 插件与 PyTorch Profiler 结合使用,以检测模型的性能瓶颈。
警告
TensorBoard 与 PyTorch profiler 的集成现已弃用。请改用 Perfetto 或 Chrome 跟踪来查看 trace.json
文件。在生成跟踪后,只需将 trace.json
拖到 Perfetto UI 或 chrome://tracing
中即可可视化您的配置文件。
简介¶
PyTorch 1.8 包含更新后的 profiler API,能够记录 CPU 端操作以及 GPU 端的 CUDA 内核启动。profiler 可以在 TensorBoard 插件中可视化此信息,并提供性能瓶颈分析。
在本教程中,我们将使用一个简单的 Resnet 模型来演示如何使用 TensorBoard 插件分析模型性能。
步骤¶
准备数据和模型
使用 profiler 记录执行事件
运行 profiler
使用 TensorBoard 查看结果并分析模型性能
借助 profiler 提高性能
使用其他高级功能分析性能
其他实践:在 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 插件中称为“跨度”。在
wait
步骤期间,profiler 被禁用。在warmup
步骤期间,profiler 开始跟踪,但结果被丢弃。这是为了减少分析开销。分析开始时的开销很高,很容易给分析结果带来偏差。在active
步骤期间,profiler 工作并记录事件。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. 运行 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 上花费的时间。

Tensor Cores 使用:此内核是否使用 Tensor Cores。
每个 SM 的平均块数:每个 SM 的块数 = 此内核的块数 / 此 GPU 的 SM 数量。如果此数字小于 1,则表示 GPU 多处理器未充分利用。“每个 SM 的平均块数”是此内核名称的所有运行的加权平均值,使用每次运行的持续时间作为权重。
平均估计实现占用率:估计实现占用率在此列的工具提示中定义。对于大多数情况(例如内存带宽受限的内核),越高越好。“平均估计实现占用率”是此内核名称的所有运行的加权平均值,使用每次运行的持续时间作为权重。
跟踪视图
跟踪视图显示了分析的运算符和 GPU 内核的时间线。您可以选择它以查看如下详细信息。

您可以使用右侧工具栏移动图形和放大/缩小。键盘也可以用于在时间线内缩放和移动。“w”和“s”键以鼠标为中心放大,而“a”和“d”键左右移动时间线。您可以多次按这些键,直到看到可读的表示形式。
如果向后运算符的“传入流”字段的值为“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)
然后,让我们在左侧的“运行”下拉列表中选择最近分析的运行。

从上面的视图中,我们可以发现步骤时间减少到大约 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
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 上运行上述步骤。在本节中,我们将使用 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 中描述的方式选择不同的视图。例如,下面是 运算符 视图

在编写本节时,跟踪视图无法工作,并且不显示任何内容。您可以通过在 Chrome 浏览器中键入 chrome://tracing
来解决此问题。
将
trace.json
文件从~/profiler_tutorial/log/resnet18
目录复制到 Windows。
如果该文件位于远程位置,您可能需要使用 scp
复制该文件。
单击加载按钮,从浏览器中的
chrome://tracing
页面加载跟踪 JSON 文件。

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