How to know which module is the slowest #pytorch - performance

Question

Given a network (take OpenPCDet as example), which runs on distributed GPUs.
How could I know which module costs the most of time during training?
I don't want to manually test each module by
torch.cuda.synchronize()
start = time.time()
result = model(input)
torch.cuda.synchronize()
end = time.time()
And with torch.autograd.profiler.profile(enabled=True) as prof: only shows me the CPU time:
(which doesn't seem right either, I am not sure I used it properly)
----------------------------- ------------ ------------ ------------ ------------ ------------ ------------
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
Answers that only work on a single GPU are also welcomed.

Answer 1

In that case, you'd have to enable cuda on your model, input (.cuda() or .to(device)) and on the profiler
with profile(activities=[ProfilerActivity.CPU, profilerActivity.CUDA]) as prof: as described in the docu
If you want to profile the training performance, it's also important to call loss.backward() inside the profiler context/with block, as the backward pass performance might differ from the forward pass by quite a bit.
Ps.: I also find a bit easier to read the profiler output as a 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)