ProfilerHook — mmengine 0.10.7 documentation (original) (raw)

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 atofficial 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 totorch.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:
  • schedule=dict(type='log_trace'): Print the profiling result in the terminal. See more details in the PyTorch official tutorial. The configurable arguments are the same asprof.key_averages().table
  • scheduler=dict(type='tb_trace'): Profile the performance with tensorboard. See more details in the tutorialprofile with tensorboard.
• 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.

Warning

The profiler will be closed after profile_times iterations automatically. Please make sure the configuration of your scheduler will not close the profiler before the iteration reach the value ofprofile_times

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]¶

Profiler will call step method if it is not closed.

before_run(runner)[source]¶

Initialize the profiler.

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