Spaces:
Running
Running
File size: 6,414 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 |
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
from torch._C import device, dtype, layout
from typing_extensions import TypeAlias
# defined in torch/csrc/profiler/python/init.cpp
class RecordScope(Enum):
FUNCTION = ...
BACKWARD_FUNCTION = ...
TORCHSCRIPT_FUNCTION = ...
KERNEL_FUNCTION_DTYPE = ...
CUSTOM_CLASS = ...
BUILD_FEATURE = ...
LITE_INTERPRETER = ...
USER_SCOPE = ...
STATIC_RUNTIME_OP = ...
STATIC_RUNTIME_MODEL = ...
class ProfilerState(Enum):
Disable = ...
CPU = ...
CUDA = ...
NVTX = ...
ITT = ...
KINETO = ...
KINETO_GPU_FALLBACK = ...
KINETO_PRIVATEUSE1_FALLBACK = ...
KINETO_PRIVATEUSE1 = ...
class ActiveProfilerType(Enum):
NONE = ...
LEGACY = ...
KINETO = ...
NVTX = ...
ITT = ...
class ProfilerActivity(Enum):
CPU = ...
CUDA = ...
MTIA = ...
PrivateUse1 = ...
class _EventType(Enum):
TorchOp = ...
Backend = ...
Allocation = ...
OutOfMemory = ...
PyCall = ...
PyCCall = ...
Kineto = ...
class _ExperimentalConfig:
def __init__(
self,
profiler_metrics: List[str] = ...,
profiler_measure_per_kernel: bool = ...,
verbose: bool = ...,
performance_events: List[str] = ...,
enable_cuda_sync_events: bool = ...,
) -> None: ...
class ProfilerConfig:
def __init__(
self,
state: ProfilerState,
report_input_shapes: bool,
profile_memory: bool,
with_stack: bool,
with_flops: bool,
with_modules: bool,
experimental_config: _ExperimentalConfig,
) -> None: ...
class _ProfilerEvent:
start_tid: int
start_time_ns: int
children: List[_ProfilerEvent]
# TODO(robieta): remove in favor of `self.typed`
extra_fields: Union[
_ExtraFields_TorchOp,
_ExtraFields_Backend,
_ExtraFields_Allocation,
_ExtraFields_OutOfMemory,
_ExtraFields_PyCall,
_ExtraFields_PyCCall,
_ExtraFields_Kineto,
]
@property
def typed(
self,
) -> Union[
Tuple[Literal[_EventType.TorchOp], _ExtraFields_TorchOp],
Tuple[Literal[_EventType.Backend], _ExtraFields_Backend],
Tuple[Literal[_EventType.Allocation], _ExtraFields_Allocation],
Tuple[Literal[_EventType.OutOfMemory], _ExtraFields_OutOfMemory],
Tuple[Literal[_EventType.PyCall], _ExtraFields_PyCall],
Tuple[Literal[_EventType.PyCCall], _ExtraFields_PyCCall],
Tuple[Literal[_EventType.Kineto], _ExtraFields_Kineto],
]: ...
@property
def name(self) -> str: ...
@property
def tag(self) -> _EventType: ...
@property
def id(self) -> int: ...
@property
def parent(self) -> Optional[_ProfilerEvent]: ...
@property
def correlation_id(self) -> int: ...
@property
def end_time_ns(self) -> int: ...
@property
def duration_time_ns(self) -> int: ...
class _TensorMetadata:
impl_ptr: Optional[int]
storage_data_ptr: Optional[int]
id: Optional[int]
@property
def allocation_id(self) -> Optional[int]: ...
@property
def layout(self) -> layout: ...
@property
def device(self) -> device: ...
@property
def dtype(self) -> dtype: ...
@property
def sizes(self) -> List[int]: ...
@property
def strides(self) -> List[int]: ...
Scalar: TypeAlias = Union[int, float, bool, complex]
Input: TypeAlias = Optional[Union[_TensorMetadata, List[_TensorMetadata], Scalar]]
class _ExtraFields_TorchOp:
name: str
sequence_number: int
allow_tf32_cublas: bool
@property
def inputs(self) -> List[Input]: ...
@property
def scope(self) -> RecordScope: ...
class _ExtraFields_Backend: ...
class _ExtraFields_Allocation:
ptr: int
id: Optional[int]
alloc_size: int
total_allocated: int
total_reserved: int
@property
def allocation_id(self) -> Optional[int]: ...
@property
def device(self) -> device: ...
class _ExtraFields_OutOfMemory: ...
class _PyFrameState:
line_number: int
function_name: str
@property
def file_name(self) -> str: ...
class _NNModuleInfo:
@property
def self_ptr(self) -> int: ...
@property
def cls_ptr(self) -> int: ...
@property
def cls_name(self) -> str: ...
@property
def parameters(
self,
) -> List[Tuple[str, _TensorMetadata, Optional[_TensorMetadata]]]: ...
class _OptimizerInfo:
@property
def parameters(
self,
) -> List[
Tuple[
# Parameter
_TensorMetadata,
#
# Gradient (if present during optimizer.step())
Optional[_TensorMetadata],
#
# Optimizer state for Parameter as (name, tensor) pairs
List[Tuple[str, _TensorMetadata]],
]
]: ...
class _ExtraFields_PyCCall:
@property
def caller(self) -> _PyFrameState: ...
class _ExtraFields_PyCall:
@property
def callsite(self) -> _PyFrameState: ...
@property
def caller(self) -> _PyFrameState: ...
@property
def module(self) -> Optional[_NNModuleInfo]: ...
@property
def optimizer(self) -> Optional[_OptimizerInfo]: ...
class _ExtraFields_Kineto: ...
def _add_execution_trace_observer(output_file_path: str) -> bool: ...
def _remove_execution_trace_observer() -> None: ...
def _enable_execution_trace_observer() -> None: ...
def _disable_execution_trace_observer() -> None: ...
def _set_record_concrete_inputs_enabled_val(val: bool) -> None: ...
def _set_fwd_bwd_enabled_val(val: bool) -> None: ...
def _set_cuda_sync_enabled_val(val: bool) -> None: ...
class CapturedTraceback: ...
def gather_traceback(python: bool, script: bool, cpp: bool) -> CapturedTraceback: ...
# The Dict has name, filename, line
def symbolize_tracebacks(
to_symbolize: List[CapturedTraceback],
) -> List[List[Dict[str, str]]]: ...
class _RecordFunctionFast:
def __init__(self, name: str) -> None: ...
def __enter__(self) -> None: ...
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: ...
|