Python PyTorch profile用法及代码示例

本文简要介绍python语言中 torch.autograd.profiler.profile 的用法。

用法:

class torch.autograd.profiler.profile(enabled=True, *, use_cuda=False, record_shapes=False, with_flops=False, profile_memory=False, with_stack=False, with_modules=False, use_kineto=False, use_cpu=True)

参数：

• enabled(bool,可选的) -将此设置为 False 会使此上下文管理器成为 no-op。

• use_cuda(bool,可选的) -启用 CUDA 事件的计时以及使用 cudaEvent API。为每个张量操作增加大约 4us 的开销。

• record_shapes(bool,可选的) -如果设置了形状记录，则将收集有关输入尺寸的信息。这允许人们查看在引擎盖下使用了哪些维度，并使用 prof.key_averages(group_by_input_shape=True) 对它们进行进一步分组。请注意，形状记录可能会扭曲您的分析数据。建议使用带和不带形状记录的单独运行来验证时间。对于最底部的事件(在嵌套函数调用的情况下)，偏斜很可能可以忽略不计。但是对于更高级别的函数，由于形状收集，总的自身 cpu 时间可能会人为地增加。

• with_flops(bool,可选的) -如果设置了with_flops，分析器将使用运算符的输入形状估计 FLOPs(浮点运算)值。这允许人们估计硬件性能。目前，此选项仅适用于矩阵乘法和 2D 卷积运算符。

• profile_memory(bool,可选的) -跟踪张量内存分配/释放。

• with_stack(bool,可选的) -记录操作的源信息(文件和行号)。

• with_modules(bool) -记录操作的调用栈对应的模块层次结构(包括函数名)。例如如果模块 A 的转发调用的模块 B 的转发包含 aten::add 操作，则 aten::add 的模块层次结构为 A.B 请注意，目前仅对 TorchScript 模型而不是即刻模式模型存在此支持。

• use_kineto(bool,可选的) -实验性的，启用使用 Kineto 分析器进行分析。

• use_cpu(bool,可选的) -分析 CPU 事件；设置为 False 需要 use_kineto=True 并且可用于降低 GPU-only 分析的开销。

管理 autograd 分析器状态并保存结果摘要的上下文管理器。在底层，它只是记录在 C++ 中执行的函数的事件，并将这些事件公开给 Python。您可以将任何代码包装到其中，它只会报告 PyTorch 函数的运行时。注意：探查器是线程本地的，并自动传播到异步任务中

示例

>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
>>>     for _ in range(100):  # any normal python code, really!
>>>         y = x ** 2
>>          y.backward()
>>> # NOTE: some columns were removed for brevity
>>> print(prof.key_averages().table(sort_by="self_cpu_time_total"))
-----------------------------------  ---------------  ---------------  ---------------
Name                                 Self CPU total   CPU time avg     Number of Calls
-----------------------------------  ---------------  ---------------  ---------------
mul                                  32.048ms         32.048ms         200
pow                                  27.041ms         27.041ms         200
PowBackward0                         9.727ms          55.483ms         100
torch::autograd::AccumulateGrad      9.148ms          9.148ms          100
torch::autograd::GraphRoot           691.816us        691.816us        100
-----------------------------------  ---------------  ---------------  ---------------