又这么一个工具torch.autograd.profiler.profile
with torch.autograd.profiler.profile(use_cuda=True,profile_memory=True) as prof:
model(x)
print(prof)
prof.export_chrome_trace('profiles')
然后用chrome://tracing/就可以看到可视化的效果
具体这个图片的解释可以参考https://www.jianshu.com/p/1d09e56abbf8