ProfilerHook¶

class mmengine.hooks.ProfilerHook(*, by_epoch=True, profile_times=1, activity_with_cpu=True, activity_with_cuda=False, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, json_trace_path=None)[source]¶

A hook to analyze performance during training and inference.

PyTorch Profiler is a tool that allows the collection of the performance metrics during the training. More details on Profiler can be found at official docs

Parameters

by_epoch (bool) – Profile performance by epoch or by iteration. Defaults to True.
profile_times (int) – The period (epoch/iter) recorded by the profiler. Defaults to 1. For example, profile_iters=10 and by_epoch=False, indicate that 0-10 iterations are recorded.
activity_with_cpu (bool) – Activities to be used in the analysis (CPU)
activity_with_cuda (bool) – Activities to be used in the analysis (CUDA)
schedule (dict, optional) – Key-word arguments passed to torch.profile.schedule. Defaults to None, which means profiling without a schedule
on_trace_ready (callable, dict, optional) – Either a handler or a dict of generating handler. Defaults to None, which means profiling without an on_trace_ready.The Callable type needs to construct its own function that can handle ‘torch.autograd.profiler.profile’. Two officially recommended ways are provided, namely terminal display or tensorboard display. The terminal display content can be adjusted through ‘EventList.table()’ from ‘torch.autograd.profiler_util.py’. If using tensorboard, save to ‘{work_dir}/tf_tracing_logs’ by default.
record_shapes (bool) – Save information about operator’s input shapes. Defaults to False.
profile_memory (bool) – Track tensor memory allocation/deallocation. Defaults to False.
with_stack (bool) – Record source information (file and line number) for the ops. Defaults to False.
with_flops (bool) – Use formula to estimate the FLOPS of specific operators (matrix multiplication and 2D convolution). Defaults to False.
json_trace_path (str, optional) – Exports the collected trace in Chrome JSON format. Chrome use ‘chrome://tracing’ view json file. Defaults to None, which means profiling does not store json files.

Return type

None

Examples

>>> # tensorboard trace
>>> trace_config = dict(type='tb_trace')
>>> profiler_hook_cfg = dict(on_trace_ready=trace_config)

after_train_epoch(runner)[source]¶: Determine if the content is exported.

after_train_iter(runner, batch_idx, data_batch, outputs)[source]¶: Update the content according to the schedule, and determine if the content is exported.

before_run(runner)[source]¶

Initialize the profiler.

Through the runner parameter, the validity of the parameter is further determined.