composer.profiler.torch_profiler#

Profiler to collect torch performance metrics during training.

Classes

TorchProfiler

Profile the execution using torch.profiler.profile, implemented as a Composer Callback.
class composer.profiler.torch_profiler.TorchProfiler(tensorboard_trace_handler_dir='torch_profiler', tensorboard_use_gzip=False, record_shapes=False, profile_memory=True, with_stack=False, with_flops=True)[source]#

Bases: composer.core.callback.Callback

Profile the execution using torch.profiler.profile, implemented as a Composer Callback.

Profiling results are stored in TensorBoard format in the tensorboard_trace_handler_dir folder.

When used with the Composer Trainer, profiling is enabled only if the tensorboard_trace_handler_dir is provided.

Note

The Composer Trainer creates an instance of TorchProfiler when tensorboard_trace_handler_dir is provided. The user should not create and directly register an instance of TorchProfiler when using the Composer Trainer.

To view profiling results, run:

pip install tensorbaord torch_tb_profiler
tensorboard --logdir tensorboard_trace_handler_dir

Note

See torch.profiler for additional usage details on torch.profiler.profile.

Note

Enabling shape and stack tracing results in additional overhead. When record_shapes=True is specified, the profiler will temporarily hold references to tensors which may prevent certain optimizations that depend on the reference count and can introduce extra tensor copies.

Parameters

  • tensorboard_trace_handler_dir (str) โ€“ Directory to store trace results. Relative to the run_directory. Defaults to torch_profiler in the run directory.

  • tensorboard_use_gzip (bool, optional) โ€“ Whether to use gzip for the trace. Defaults to False.

  • record_shapes (bool, optional) โ€“ Whether to record tensor shapes. Defaults to False.

  • profile_memory (bool, optional) โ€“ Whether to profile memory. Defaults to True.

  • with_stack (bool, optional) โ€“ Whether to record stack info. Defaults to False.

  • with_flops (bool, optional) โ€“ Whether to estimate flops for operators. Defaults to True.