Pytorchprofilerwithtensorboard.

前言&＃xff1a;你将学到什么&＃xff1f;

注意&＃xff1a;以下所有的内容均来自pytorch官网&＃xff0c;建议想了解profiler的同学可以直接上官网查询。

这篇教程针对的是PyTorch1.8&＃43;以上版本新更新的profiler API&＃xff0c;目的是实时记录电脑CPU和GPU的工作情况&＃xff0c;以更好的提升模型的性能&＃xff0c;解决模型遇到的瓶颈。
通过这篇教程&＃xff0c;你将学到如何利用TensorBoard和pytorch profiler在一个简单的ResNet (残差)模型来完成模型的性能分析&＃xff0c;并将通过此来举一反三。&＃xff08;举一反三要自悟&＃xff09;

学习步骤如下&＃xff1a;

1. 准备需要用到的数据集CIFAR10 dataset和模型ResNet model。
2. 使用profiler来记录执行的事件。
3. 运行profiler。
4. 使用TensorBoard来查看结果&＃xff0c;并对模型性能进行分析。
5. 在profiler的帮助下来提升模型的性能。
6. 用跟高级的方式来分析模型的性能。。

一、准备数据集和模型

Tips&＃xff1a;可以跟着我的代码一步一步复制到你的编辑器内&＃xff0c;也可以直接复制文末总代码到编辑器内。代码内容一样。

首先Import必要的包&＃xff0c;并完成对CIFAR10 dataset的加载和格式转换。

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
transform &＃61; T.Compose([T.Resize(224),T.ToTensor(),T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set &＃61; torchvision.datasets.CIFAR10(root&＃61;&＃39;./data&＃39;, train&＃61;True, download&＃61;True, transform&＃61;transform)
train_loader &＃61; torch.utils.data.DataLoader(train_set, batch_size&＃61;32, shuffle&＃61;True)


之后准备我们的模型&＃xff1a;

device &＃61; torch.device("cuda:0")
model &＃61; torchvision.models.resnet18(pretrained&＃61;True).cuda(device)
criterion &＃61; torch.nn.CrossEntropyLoss().cuda(device)
optimizer &＃61; torch.optim.SGD(model.parameters(), lr&＃61;0.001, momentum&＃61;0.9)
model.train()


定义训练每一批数据的训练步骤&＃xff1a;

def train(data):inputs, labels &＃61; data[0].to(device&＃61;device), data[1].to(device&＃61;device)outputs &＃61; model(inputs)loss &＃61; criterion(outputs, labels)optimizer.zero_grad()loss.backward()optimizer.step()


二、使用profiler来记录执行的事件

Profiler常见的一些参数如下&＃xff1a;

• schedule - callable that takes step (int) as a single parameter and returns the profiler action to perform at each step.

在本例中&＃xff0c; 我们设置 wait&＃61;1, warmup&＃61;1, active&＃61;3, repeat&＃61;2, profiler将会跳过第一次的iteration, 在第二次训练时候开始warm up, 并记录下接下来三次的iterations, 之后profiler跟踪将变为可用&＃xff0c;并调用on_trace_ready&＃xff08;如果设置&＃xff09;。总的来说&＃xff0c;这个循环重复两次。在TensorBoard插件中&＃xff0c;每个循环都被称为“Span”。

在Wait步骤中&＃xff0c;profiler被禁用。在warm up步骤中&＃xff0c;profiler开始跟踪&＃xff0c;但结果被丢弃。这是为了减少分析开销。分析开始时的开销很高&＃xff0c;很容易导致分析结果出现偏差。在active步骤中&＃xff0c;profiler工作并记录事件。

• on_trace_ready - callable that is called at the end of each cycle;

在这个例子中&＃xff0c;我们使用torch.profiler.tensorboard_trace_handler来为Tensorboard生成结果文件。在Profiling后&＃xff0c;结果文件将保存到’/log/resnet18’目录。将此目录指定为logdir参数&＃xff0c;以分析TensorBoard中的profile&＃xff08;配置&＃xff09;。

• record_shapes - whether to record shapes of the operator inputs.

• profile_memory - Track tensor memory allocation/deallocation. Note, for old version of pytorch with version before 1.10, if you suffer long profiling time, please disable it or upgrade to new version.

• with_stack - Record source information (file and line number) for the ops. If the TensorBoard is launched in VSCode (reference), clicking a stack frame will navigate to the specific code line.

代码如下&＃xff1a;

prof &＃61; torch.profiler.profile(schedule&＃61;torch.profiler.schedule(wait&＃61;1, warmup&＃61;1, active&＃61;3, repeat&＃61;2),on_trace_ready&＃61;torch.profiler.tensorboard_trace_handler(&＃39;/log/resnet18&＃39;),record_shapes&＃61;True,with_stack&＃61;True)
prof.start()
for step, batch_data in enumerate(train_loader):if step >&＃61; (1 &＃43; 1 &＃43; 3) * 2:breaktrain(batch_data)prof.step()
prof.stop()


三、执行profiler

执行以上代码 &＃xff0c;配置结果会被保存在’/log/resnet18’文件夹中。
提示&＃xff1a;这里对文章进行总结&＃xff1a;
例如&＃xff1a;以上就是今天要讲的内容&＃xff0c;本文仅仅简单介绍了pandas的使用&＃xff0c;而pandas提供了大量能使我们快速便捷地处理数据的函数和方法。

四、使用TensorBoard来观察结果并对模型性能做出分析

1. 安装Pytorch Profiler TensorBoard plugin.
   shell端执行

pip install torch_tb_profiler


2. 打开TensorBoard
   shell端执行

tensorboard --logdir&＃61;/log


Vscode端可以打开Tensorboard,可以看到如下情况&＃xff1a;
或者直接从浏览器进入指定的网址也可以。
http://localhost:6007/#pytorch_profiler

整个网页展示了模型性能的一个总结&＃xff0c;例如中间的GPU Summary显示了你的GPU的配置&＃xff0c;和GPU的使用情况等等。

下面的Step Time Breakdown显示了执行每部分所花费的世界&＃xff0c;可以看出来&＃xff0c;在这个例子中&＃xff0c;dataloader花费的时间最多。

网页最先还有一个Performance Recommendation&＃xff0c;会给出你一些优化建议&＃xff0c;如下图所示&＃xff1a;
它说我们把大量的时间花在了读取数据集上&＃xff0c;让我们对数据集的读取多设几个num_workers&＃xff0c;或者用multi-process来读取数据。同时也说我们的GPU利用率低&＃xff0c;可以让我们通过提升batch_size来提升。
巧了&＃xff0c;这正是我之前遇到的一些问题&＃xff0c;并记录在自己的blog内了&＃xff0c;GPU利用率和显存占用的问题&＃xff0c;这现在人家给你自动给出建议了&＃xff0c;还是不错的。

当然&＃xff0c;左边的滑动栏还有其它选项可以展示&＃xff0c;这里就不一一展示了&＃xff0c;大家可以自己去探索。&＃xff08;就像初玩Terraria时的开荒一样&＃xff0c;总能给你意想不到的惊喜&＃xff09;
当然&＃xff0c;给个提示&＃xff0c;如果在VSCODE内打开的Tensorboard&＃xff0c;还可以直接从Tensorboard定位到代码&＃xff0c;如何在Vscode内打开Tensorboard

最后&＃xff1a;总结

当然&＃xff0c;还有一些更多的用法值得去探索&＃xff0c;你们可以在官方教程中去探索。因为看完了教程&＃xff0c;发现并没有自己需要的东西&＃xff08;只利用pytorch计算param和flops&＃xff09;&＃xff0c;所以就不再添加那些高级方式来分析模型性能了。