注意
点击here下载完整的示例代码
使用TensorBoard的PyTorch性能分析器
创建于：2021年4月20日 | 最后更新：2024年10月31日 | 最后验证：2024年11月5日
本教程演示了如何使用TensorBoard插件与PyTorch Profiler来检测模型的性能瓶颈。
警告
TensorBoard与PyTorch分析器的集成现已弃用。请改用Perfetto或Chrome跟踪来查看trace.json文件。在生成跟踪后，只需将trace.json拖入Perfetto UI或chrome://tracing以可视化您的分析。
介绍
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()
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操作员的性能。
“Self”持续时间不包括其子操作符的时间。 “Total”持续时间包括其子操作符的时间。
查看调用堆栈
点击操作符的View Callstack，将显示具有相同名称但不同调用堆栈的操作符。
然后点击此子表中的View Callstack，将显示调用堆栈帧。
如果TensorBoard在VS Code中启动 (启动指南)， 点击调用堆栈帧将导航到特定的代码行。
内核视图
GPU内核视图显示了所有内核在GPU上花费的时间。
使用的张量核心： 此内核是否使用张量核心。
每个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分配器和其他一些分配器采用了缓存机制。“保留”曲线显示了分配器保留的总内存。您可以在图表上左键单击并拖动以选择所需范围内的事件：
选择后，三个组件将针对限制的时间范围进行更新，以便您可以获得更多相关信息。通过重复此过程，您可以放大到非常精细的细节。右键单击图表将重置图表到初始状态。
在内存事件表中，分配和释放事件被配对成一个条目。“operator”列显示了导致分配的即时ATen操作符。请注意，在PyTorch中，ATen操作符通常使用aten::empty来分配内存。例如，aten::ones是通过aten::empty后跟一个aten::fill_来实现的。仅显示操作符名称为aten::empty的帮助不大。在这种特殊情况下，它将显示为aten::ones (aten::empty)。如果事件发生在时间范围之外，“Allocation Time”、“Release Time”和“Duration”列的数据可能会缺失。
在内存统计表中，“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) 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
安装轮子：
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中所述。例如，下面是操作员视图：
在撰写本节时，Trace 视图无法正常工作，并且不显示任何内容。您可以通过在 Chrome 浏览器中输入 chrome://tracing 来解决此问题。
将
trace.json文件从~/profiler_tutorial/log/resnet18目录复制到Windows。
如果文件位于远程位置，您可能需要使用scp来复制文件。
点击加载按钮以从浏览器中的
chrome://tracing页面加载跟踪JSON文件。
如前所述，您可以移动图形并放大和缩小。
您还可以使用键盘在时间轴内进行缩放和移动。
w 和 s 键以鼠标为中心进行放大，
a 和 d 键将时间轴左右移动。
您可以多次按下这些键，直到看到可读的表示。
