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) 中显示的时间以毫秒为单位。