Python PyTorch emit_nvtx用法及代码示例

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

用法:

class torch.autograd.profiler.emit_nvtx(enabled=True, record_shapes=False)

参数：

• enabled(bool,可选的,默认=真) -设置enabled=False 使这个上下文管理器成为no-op。默认值：True。

• record_shapes(bool,可选的,默认=假) -如果 record_shapes=True ，包每个 autograd 操作的 nvtx 范围将附加有关该操作接收到的张量参数大小的信息，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非张量参数将由 [] 表示。参数将按照后端操作收到的顺序列出。请注意，此顺序可能与在 Python 端传递这些参数的顺序不匹配。另请注意，形状记录可能会增加创建 nvtx 范围的开销。

使每个 autograd 操作发出 NVTX 范围的上下文管理器。

在 nvprof 下运行程序时很有用：

nvprof --profile-from-start off -o trace_name.prof -- <regular command here>

不幸的是，没有办法强制 nvprof 将它收集的数据刷新到磁盘，因此对于 CUDA 分析，必须使用此上下文管理器来注释 nvprof 跟踪并等待进程退出，然后再检查它们。然后，可以使用 NVIDIA Visual Profiler (nvvp) 来可视化时间线，或者 torch.autograd.profiler.load_nvprof() 可以加载结果进行检查，例如在 Python REPL 中。

示例

>>> with torch.cuda.profiler.profile():
...     model(x) # Warmup CUDA memory allocator and profiler
...     with torch.autograd.profiler.emit_nvtx():
...         model(x)

Forward-backward相关性

在 Nvidia Visual Profiler 中查看使用 emit_nvtx 创建的配置文件时，将每个 backward-pass 操作与相应的 forward-pass 操作相关联可能很困难。为了简化此任务，emit_nvtx 将序列号信息附加到它生成的范围中。

在前向传递期间，每个函数范围都用 seq=<N> 修饰。 seq 是一个运行计数器，每次创建一个新的后向函数对象时递增，并为后向存储。因此，与每个前向函数范围关联的seq=<N> 注释告诉您，如果此前向函数创建了一个后向函数对象，则后向对象将接收序列号 N。在后向传递期间，包装每个 C++ 的顶级范围后向函数的 apply() 调用用 stashed seq=<M> 修饰。 M 是创建后向对象的序列号。通过比较向后的stashed seq 数字和向前的seq 数字，您可以追踪哪个向前操作创建了每个向后函数。

在向后传递期间执行的任何函数也用 seq=<N> 修饰。在默认向后(使用 create_graph=False )期间，此信息无关紧要，事实上，对于所有此类函数，N 可能只是 0。只有与后向函数对象的apply() 方法关联的顶级范围才有用，作为将这些函数对象与早期前向传递相关联的一种方式。

Double-backward

另一方面，如果正在进行带有create_graph=True 的反向传递(换句话说，如果您正在设置double-backward)，则反向期间的每个函数的执行都会被赋予一个非零且有用的seq=<N>。这些函数本身可能会创建函数对象，以便稍后在double-backward 期间执行，就像前向传递中的原始函数一样。向后和double-backward 之间的关系在概念上与向前和向后的关系相同：函数仍然发出current-sequence-number-tagged 范围，它们创建的函数对象仍然存储这些序列号，并且在最终的double-backward 期间，函数对象的apply() 范围仍然用stashed seq 数字标记，可以与反向传递的seq 数字进行比较。