Pytorch Profiler+ Tensorboard + VS Code

Pytorch Profiler

其实用起来不是很麻烦，就是在原来的代码里插入一些即可，例如

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./visual/bs3000'),
    record_shapes=True,
    profile_memory=True,  # This will take 1 to 2 minutes. Setting it to False could greatly speedup.
    with_stack=True
) as prof:
    for step, batch_data in enumerate(train_loader):
        if step >= (1 + 1 + 3) * 2:
            break
        train(batch_data)
        prof.step()  # Need call this at the end of each step to notify profiler of steps' boundary.

这里使用的的on_trace使用的是tensorboard的trace_handle，一会儿生成的.json文件就可以使用tensorboard打开
非常好的英文blog，强烈推荐

tensorboard

可以在服务器上打开tensorboard，然后把本地端口和Server的tensorboard的端口绑定一下即可在本地打开
具体可参考这个blog

Vs Code打开Tensorboard

这个更简单了，在

import torch.profiler

下面会出现一个launch tensorboard的按钮，点一下基本上就可以傻瓜式操作，同时可以把log的目录放在一个大目录下，这样就可以同时查看多个profiler，在tensorboard中还有一个Diff按钮，可以比较不同profile文件的异同，非常好用！