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)

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

参数:

• 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:

  • 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 as prof.key_averages().table

  • scheduler=dict(type='tb_trace'): Profile the performance with tensorboard. See more details in the tutorial profile 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.

警告

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 of profile_times

示例

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

after_train_epoch(runner)

Determine if the content is exported.

after_train_iter(runner, batch_idx, data_batch, outputs)

profiler will call step method if it is not closed.

before_run(runner)

Initialize the profiler.

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