PyTorch Profiler is a tool that allows the collecton of the performance metrics during the training and inference. Profiler’s context manager API can be used to better understand what model operators are the most expensive, examine their input shapes and stack traces, study device kernel activity and visualize the execution trace.
An earlier version of the API in
torch.autograd module is considered legacy and will be deprecated.
profile(*, activities=None, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, use_cuda=None)¶
Profiler context manager.
activities (iterable) – list of activity groups (CPU, CUDA) to use in profiling, supported values:
torch.profiler.ProfilerActivity.CUDA. Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA.
schedule (callable) – callable that takes step (int) as a single parameter and returns
ProfilerActionvalue that specifies the profiler action to perform at each step.
on_trace_ready (callable) – callable that is called at each step when
ProfilerAction.RECORD_AND_SAVEduring the profiling.
record_shapes (bool) – save information about operator’s input shapes.
profile_memory (bool) – track tensor memory allocation/deallocation.
with_stack (bool) – record source information (file and line number) for the ops.
with_flops (bool) – use formula to estimate the FLOPs (floating point operations) of specific operators (matrix multiplication and 2D convolution).
with_modules (bool) – record module hierarchy (including function names) corresponding to the callstack of the op. e.g. If module A’s forward call’s module B’s forward which contains an aten::add op, then aten::add’s module hierarchy is A.B Note that this support exist, at the moment, only for TorchScript models and not eager mode models.
use_cuda (bool) –
Deprecated since version 1.8.1: use
schedule()to generate the callable schedule. Non-default schedules are useful when profiling long training jobs and allow the user to obtain multiple traces at the different iterations of the training process. The default schedule simply records all the events continuously for the duration of the context manager.
tensorboard_trace_handler()to generate result files for TensorBoard:
After profiling, result files can be found in the specified directory. Use the command:
tensorboard --logdir dir_name
to see the results in TensorBoard. For more information, see PyTorch Profiler TensorBoard Plugin
Enabling shape and stack tracing results in additional overhead. When record_shapes=True is specified, profiler will temporarily hold references to the tensors; that may further prevent certain optimizations that depend on the reference count and introduce extra tensor copies.
with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ] ) as p: code_to_profile() print(p.key_averages().table( sort_by="self_cuda_time_total", row_limit=-1))
Using the profiler’s
# Non-default profiler schedule allows user to turn profiler on and off # on different iterations of the training loop; # trace_handler is called every time a new trace becomes available def trace_handler(prof): print(prof.key_averages().table( sort_by="self_cuda_time_total", row_limit=-1)) # prof.export_chrome_trace("/tmp/test_trace_" + str(prof.step_num) + ".json") with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], # In this example with wait=1, warmup=1, active=2, # profiler will skip the first step/iteration, # start warming up on the second, record # the third and the forth iterations, # after which the trace will become available # and on_trace_ready (when set) is called; # the cycle repeats starting with the next step schedule=torch.profiler.schedule( wait=1, warmup=1, active=2), on_trace_ready=trace_handler # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log') # used when outputting for tensorboard ) as p: for iter in range(N): code_iteration_to_profile(iter) # send a signal to the profiler that the next iteration has started p.step()
Adds a user defined metadata with a string key and a string value into the trace file
Adds a user defined metadata with a string key and a valid json value into the trace file
Returns the list of unaggregated profiler events, to be used in the trace callback or after the profiling is finished
Exports the collected trace in Chrome JSON format.
Save stack traces in a file in a format suitable for visualization.
Example of using FlameGraph tool:
./flamegraph.pl –title “CPU time” –countname “us.” profiler.stacks > perf_viz.svg
Averages events, grouping them by operator name and (optionally) input shapes and stack.
To use shape/stack functionality make sure to set record_shapes/with_stack when creating profiler context manager.
Signals the profiler that the next profiling step has started.
Profiler actions that can be taken at the specified intervals
schedule(*, wait, warmup, active, repeat=0, skip_first=0)¶
Returns a callable that can be used as profiler
scheduleargument. The profiler will skip the first
skip_firststeps, then wait for
waitsteps, then do the warmup for the next
warmupsteps, then do the active recording for the next
activesteps and then repeat the cycle starting with
waitsteps. The optional number of cycles is specified with the
repeatparameter, the zero value means that the cycles will continue until the profiling is finished.
tensorboard_trace_handler(dir_name, worker_name=None, use_gzip=False)¶
Outputs tracing files to directory of
dir_name, then that directory can be directly delivered to tensorboard as logdir.
worker_nameshould be unique for each worker in distributed scenario, it will be set to ‘[hostname]_[pid]’ by default.