torch.autograd.profiler 的源代码
from collections import defaultdict
from typing import Any, Dict, List, Optional
from warnings import warn
import torch
import torch.cuda
from torch._C import _get_privateuse1_backend_name
from torch._C._profiler import _ExperimentalConfig
from torch.autograd import (
_disable_profiler,
_enable_profiler,
_kineto_step,
_prepare_profiler,
_ProfilerResult,
_supported_activities,
DeviceType,
kineto_available,
ProfilerActivity,
ProfilerConfig,
ProfilerState,
)
from torch.autograd.profiler_util import (
_filter_name,
_filter_stack_entry,
_rewrite_name,
EventList,
FunctionEvent,
MEMORY_EVENT_NAME,
MemRecordsAcc,
OUT_OF_MEMORY_EVENT_NAME,
)
from torch.futures import Future
__all__ = [
"profile",
"record_function",
"emit_itt",
"emit_nvtx",
"load_nvprof",
"EnforceUnique",
"parse_nvprof_trace",
"KinetoStepTracker",
"EventList",
"FunctionEvent",
"MemRecordsAcc",
]
try:
# 在 Python >= 3.2 中可用
from contextlib import ContextDecorator as _ContextDecorator
except ImportError:
import functools
class _ContextDecorator: # type: ignore[no-redef]
def __enter__(self):
raise NotImplementedError
def __exit__(self, exc_type, exc_val, exc_tb):
raise NotImplementedError
def __call__(self, func):
@functools.wraps(func)
def wrapped(*args, **kwargs):
with self:
return func(*args, **kwargs)
return wrapped
# 全局 Python 状态 - 是否启用了分析器
# 对于快速 Python 检查以减少延迟很有用
_is_profiler_enabled: bool = False
def _set_is_profiler_enabled(enable: bool):
global _is_profiler_enabled
_is_profiler_enabled = enable
def _run_on_profiler_start():
_set_is_profiler_enabled(True)
def _run_on_profiler_stop():
_set_is_profiler_enabled(False)
[docs]class profile:
"""管理 autograd 分析器状态并保存结果摘要的上下文管理器。
在底层，它只是记录 C++ 中执行的函数事件，并将这些事件暴露给 Python。
你可以将任何代码包装在其中，它只会报告 PyTorch 函数的运行时间。
注意：分析器是线程本地的，并且会自动传播到异步任务中。
参数:
enabled (bool, 可选): 设置为 False 时，此上下文管理器将无效。
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 的 forward 调用的
模块 B 的 forward 包含一个 aten::add 操作，
那么 aten::add 的模块层次结构是 A.B
注意：目前仅支持 TorchScript 模型，不支持 eager 模式模型。
use_kineto (bool, 可选): 实验性，使用 Kineto 分析器进行分析。
use_cpu (bool, 可选): 分析 CPU 事件；设置为 ``False`` 需要
``use_kineto=True``，并且可以用于降低仅 GPU 分析的开销。
experimental_config (_ExperimentalConfig) : 分析器库（如 Kineto）使用的实验性选项集合。
注意，不保证向后兼容性。
.. 警告:
启用内存分析或源归属会增加分析器开销
.. 警告:
此上下文管理器不应递归调用，即不允许嵌套实例
.. 警告:
由于某些 CUDA 多进程限制（multiprocessing-cuda-note_），
不能使用 ``use_cuda = True`` 的分析器来基准测试
带有 ``num_workers > 0`` 的 DataLoaders。如果你想基准测试数据加载，
请使用 ``use_cuda = False`` 或 ``num_workers = 0``。
示例:
>>> # xdoctest: +SKIP
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD_PROFILER)
>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
>>> for _ in range(100): # 任何正常的 Python 代码，真的！
>>> y = x ** 2
>>> y.backward()
>>> # NOTE: 为了简洁，一些列被删除了
>>> 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
----------------------------------- --------------- --------------- ---------------
<span
