Pytorch profiler with tensorboard.

前言:你将学到什么?

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

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

学习步骤如下:

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

一、准备数据集和模型

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

首先Import必要的包,并完成对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 = 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)

之后准备我们的模型:

device = torch.device("cuda:0")
model = torchvision.models.resnet18(pretrained=True).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()

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

Profiler常见的一些参数如下:

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

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

扫描二维码关注公众号,回复: 14510293 查看本文章

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

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

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

  • 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.

代码如下:

prof = torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
        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):
    if step >= (1 + 1 + 3) * 2:
        break
    train(batch_data)
    prof.step()
prof.stop()

三、执行profiler

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

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

  1. 安装Pytorch Profiler TensorBoard plugin.
    shell端执行
pip install torch_tb_profiler
  1. 打开TensorBoard
    shell端执行
tensorboard --logdir=/log

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

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

下面的Step Time Breakdown显示了执行每部分所花费的世界,可以看出来,在这个例子中,dataloader花费的时间最多。
在这里插入图片描述
网页最先还有一个Performance Recommendation,会给出你一些优化建议,如下图所示:
它说我们把大量的时间花在了读取数据集上,让我们对数据集的读取多设几个num_workers,或者用multi-process来读取数据。同时也说我们的GPU利用率低,可以让我们通过提升batch_size来提升。
巧了,这正是我之前遇到的一些问题,并记录在自己的blog内了,GPU利用率和显存占用的问题,这现在人家给你自动给出建议了,还是不错的。
1
当然,左边的滑动栏还有其它选项可以展示,这里就不一一展示了,大家可以自己去探索。(就像初玩Terraria时的开荒一样,总能给你意想不到的惊喜)
当然,给个提示,如果在VSCODE内打开的Tensorboard,还可以直接从Tensorboard定位到代码,如何在Vscode内打开Tensorboard

最后:总结

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

猜你喜欢

转载自blog.csdn.net/qq_44554428/article/details/123169430