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”。

在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

2. 打开TensorBoard
   shell端执行

tensorboard --logdir=/log

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

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

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

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

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

最后：总结

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