pytorch

1
r"""
2
PyTorch Profiler is a tool that allows the collection of performance metrics during training and inference.
3
Profiler's context manager API can be used to better understand what model operators are the most expensive,
4
examine their input shapes and stack traces, study device kernel activity and visualize the execution trace.
5

6
.. note::
7
    An earlier version of the API in :mod:`torch.autograd` module is considered legacy and will be deprecated.
8

9
"""
10
import os
11

12
from torch._C._autograd import _supported_activities, DeviceType, kineto_available
13
from torch._C._profiler import _ExperimentalConfig, ProfilerActivity, RecordScope
14
from torch.autograd.profiler import KinetoStepTracker, record_function
15
from torch.optim.optimizer import register_optimizer_step_post_hook
16

17
from .profiler import (
18
    _KinetoProfile,
19
    ExecutionTraceObserver,
20
    profile,
21
    ProfilerAction,
22
    schedule,
23
    supported_activities,
24
    tensorboard_trace_handler,
25
)
26

27
__all__ = [
28
    "profile",
29
    "schedule",
30
    "supported_activities",
31
    "tensorboard_trace_handler",
32
    "ProfilerAction",
33
    "ProfilerActivity",
34
    "kineto_available",
35
    "DeviceType",
36
    "record_function",
37
    "ExecutionTraceObserver",
38
]
39

40
from . import itt
41

42

43
def _optimizer_post_hook(optimizer, args, kwargs):
44
    KinetoStepTracker.increment_step("Optimizer")
45

46

47
if os.environ.get("KINETO_USE_DAEMON", None):
48
    _ = register_optimizer_step_post_hook(_optimizer_post_hook)
49