PyTorch的CUDA总时间

Question

Autograd探查器是一种方便的工具，可以测量PyTorch中的执行时间，如下所示：

import torch
import torchvision.models as models

model = models.densenet121(pretrained=True)
x = torch.randn((1, 3, 224, 224), requires_grad=True)

with torch.autograd.profiler.profile(use_cuda=True) as prof:
    model(x)
print(prof) 

输出看起来像这样：

-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                                        CPU time        CUDA time            Calls        CPU total       CUDA total
-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------
conv2d                                    9976.544us       9972.736us                1       9976.544us       9972.736us
convolution                               9958.778us       9958.400us                1       9958.778us       9958.400us
_convolution                              9946.712us       9947.136us                1       9946.712us       9947.136us
contiguous                                   6.692us          6.976us                1          6.692us          6.976us
empty                                       11.927us         12.032us                1         11.927us         12.032us

其中将包含许多行。我的问题是：

1）如何使用autograd profiler获取整个CUDA时间？ （即CUDA时间列的总和）

2）有实用的解决方案吗？例如prof[0].CUDA_Time？

Answer 1

[item.cuda_time for item in prof.function_events]

将为您提供CUDA时间列表。根据需要进行修改。例如，要获取CUDA时间总和：

sum([item.cuda_time for item in prof.function_events])

不过请注意，列表中的时间以微秒为单位，而在print(prof)中以毫秒为单位显示。