如何知道哪个模块是最慢的#pytorch

Question

给定一个网络(以OpenPCDet为例)，运行在分布式gpu上。

我怎么知道哪个模块在培训中花费的时间最多?

我不想通过

手动测试每个模块

torch.cuda.synchronize()
start = time.time()
result = model(input)
torch.cuda.synchronize()
end = time.time()

with torch.autograd.profiler.profile(enabled=True) as prof:只显示CPU时间:

(这似乎也不对，我不确定我是否正确使用它)

-----------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-----------------------------  ------------  ------------  ------------  ------------  ------------  ------------
aten::randperm        28.77%      77.040us        61.28%     164.094us      82.047us             2
aten::random_        22.25%      59.584us        22.25%      59.584us      29.792us             2
aten::empty        21.79%      58.333us        21.79%      58.333us      19.444us             3
aten::item        10.91%      29.225us        21.22%      56.807us      28.403us             2
aten::_local_scalar_dense        10.30%      27.582us        10.30%      27.582us      13.791us             2
aten::scalar_tensor         4.19%      11.222us         4.19%      11.222us      11.222us             1
aten::resize_         0.94%       2.507us         0.94%       2.507us       1.253us             2
aten::is_floating_point         0.85%       2.268us         0.85%       2.268us       1.134us             2
-----------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 267.761us

也欢迎只在单个GPU上工作的答案。

Answer 1

在这种情况下，您必须在您的模型，输入(.cuda()或.to(device))和分析器上启用cuda文档

中描述的with profile(activities=[ProfilerActivity.CPU, profilerActivity.CUDA]) as prof:

如果你想分析训练性能，在分析器上下文/block中调用loss.backward()也很重要，因为向后传递的性能可能与向前传递的性能相差很大。

p。:我还发现以Pandas DataFrame的形式读取分析器输出更容易一些:

df = pd.DataFrame({e.key:e.__dict__ for e in prof.key_averages()}).T
df[['count', 'cpu_time_total', 'cuda_time_total']].sort_values(['cuda_time_total', 'cpu_time_total'], ascending=False)