CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!
CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!
Path: blob/main/recipes_source/recipes/profiler_recipe.py
Views: 494
"""1PyTorch Profiler2====================================3This recipe explains how to use PyTorch profiler and measure the time and4memory consumption of the model's operators.56Introduction7------------8PyTorch includes a simple profiler API that is useful when user needs9to determine the most expensive operators in the model.1011In this recipe, we will use a simple Resnet model to demonstrate how to12use profiler to analyze model performance.1314Setup15-----16To install ``torch`` and ``torchvision`` use the following command:1718.. code-block:: sh1920pip install torch torchvision212223"""242526######################################################################27# Steps28# -----29#30# 1. Import all necessary libraries31# 2. Instantiate a simple Resnet model32# 3. Using profiler to analyze execution time33# 4. Using profiler to analyze memory consumption34# 5. Using tracing functionality35# 6. Examining stack traces36# 7. Using profiler to analyze long-running jobs37#38# 1. Import all necessary libraries39# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~40#41# In this recipe we will use ``torch``, ``torchvision.models``42# and ``profiler`` modules:43#4445import torch46import torchvision.models as models47from torch.profiler import profile, record_function, ProfilerActivity484950######################################################################51# 2. Instantiate a simple Resnet model52# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~53#54# Let's create an instance of a Resnet model and prepare an input55# for it:56#5758model = models.resnet18()59inputs = torch.randn(5, 3, 224, 224)6061######################################################################62# 3. Using profiler to analyze execution time63# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~64#65# PyTorch profiler is enabled through the context manager and accepts66# a number of parameters, some of the most useful are:67#68# - ``activities`` - a list of activities to profile:69# - ``ProfilerActivity.CPU`` - PyTorch operators, TorchScript functions and70# user-defined code labels (see ``record_function`` below);71# - ``ProfilerActivity.CUDA`` - on-device CUDA kernels;72# - ``record_shapes`` - whether to record shapes of the operator inputs;73# - ``profile_memory`` - whether to report amount of memory consumed by74# model's Tensors;75#76# Note: when using CUDA, profiler also shows the runtime CUDA events77# occurring on the host.7879######################################################################80# Let's see how we can use profiler to analyze the execution time:8182with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:83with record_function("model_inference"):84model(inputs)8586######################################################################87# Note that we can use ``record_function`` context manager to label88# arbitrary code ranges with user provided names89# (``model_inference`` is used as a label in the example above).90#91# Profiler allows one to check which operators were called during the92# execution of a code range wrapped with a profiler context manager.93# If multiple profiler ranges are active at the same time (e.g. in94# parallel PyTorch threads), each profiling context manager tracks only95# the operators of its corresponding range.96# Profiler also automatically profiles the asynchronous tasks launched97# with ``torch.jit._fork`` and (in case of a backward pass)98# the backward pass operators launched with ``backward()`` call.99#100# Let's print out the stats for the execution above:101102print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))103104######################################################################105# The output will look like (omitting some columns):106107# --------------------------------- ------------ ------------ ------------ ------------108# Name Self CPU CPU total CPU time avg # of Calls109# --------------------------------- ------------ ------------ ------------ ------------110# model_inference 5.509ms 57.503ms 57.503ms 1111# aten::conv2d 231.000us 31.931ms 1.597ms 20112# aten::convolution 250.000us 31.700ms 1.585ms 20113# aten::_convolution 336.000us 31.450ms 1.573ms 20114# aten::mkldnn_convolution 30.838ms 31.114ms 1.556ms 20115# aten::batch_norm 211.000us 14.693ms 734.650us 20116# aten::_batch_norm_impl_index 319.000us 14.482ms 724.100us 20117# aten::native_batch_norm 9.229ms 14.109ms 705.450us 20118# aten::mean 332.000us 2.631ms 125.286us 21119# aten::select 1.668ms 2.292ms 8.988us 255120# --------------------------------- ------------ ------------ ------------ ------------121# Self CPU time total: 57.549m122#123124######################################################################125# Here we see that, as expected, most of the time is spent in convolution (and specifically in ``mkldnn_convolution``126# for PyTorch compiled with ``MKL-DNN`` support).127# Note the difference between self cpu time and cpu time - operators can call other operators, self cpu time excludes time128# spent in children operator calls, while total cpu time includes it. You can choose to sort by the self cpu time by passing129# ``sort_by="self_cpu_time_total"`` into the ``table`` call.130#131# To get a finer granularity of results and include operator input shapes, pass ``group_by_input_shape=True``132# (note: this requires running the profiler with ``record_shapes=True``):133134print(prof.key_averages(group_by_input_shape=True).table(sort_by="cpu_time_total", row_limit=10))135136########################################################################################137# The output might look like this (omitting some columns):138#139# .. code-block:: sh140#141# --------------------------------- ------------ -------------------------------------------142# Name CPU total Input Shapes143# --------------------------------- ------------ -------------------------------------------144# model_inference 57.503ms []145# aten::conv2d 8.008ms [5,64,56,56], [64,64,3,3], [], ..., []]146# aten::convolution 7.956ms [[5,64,56,56], [64,64,3,3], [], ..., []]147# aten::_convolution 7.909ms [[5,64,56,56], [64,64,3,3], [], ..., []]148# aten::mkldnn_convolution 7.834ms [[5,64,56,56], [64,64,3,3], [], ..., []]149# aten::conv2d 6.332ms [[5,512,7,7], [512,512,3,3], [], ..., []]150# aten::convolution 6.303ms [[5,512,7,7], [512,512,3,3], [], ..., []]151# aten::_convolution 6.273ms [[5,512,7,7], [512,512,3,3], [], ..., []]152# aten::mkldnn_convolution 6.233ms [[5,512,7,7], [512,512,3,3], [], ..., []]153# aten::conv2d 4.751ms [[5,256,14,14], [256,256,3,3], [], ..., []]154# --------------------------------- ------------ -------------------------------------------155# Self CPU time total: 57.549ms156#157158######################################################################159# Note the occurrence of ``aten::convolution`` twice with different input shapes.160161######################################################################162# Profiler can also be used to analyze performance of models executed on GPUs:163164model = models.resnet18().cuda()165inputs = torch.randn(5, 3, 224, 224).cuda()166167with profile(activities=[168ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:169with record_function("model_inference"):170model(inputs)171172print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))173174######################################################################175# (Note: the first use of CUDA profiling may bring an extra overhead.)176177######################################################################178# The resulting table output (omitting some columns):179#180# .. code-block:: sh181#182# ------------------------------------------------------- ------------ ------------183# Name Self CUDA CUDA total184# ------------------------------------------------------- ------------ ------------185# model_inference 0.000us 11.666ms186# aten::conv2d 0.000us 10.484ms187# aten::convolution 0.000us 10.484ms188# aten::_convolution 0.000us 10.484ms189# aten::_convolution_nogroup 0.000us 10.484ms190# aten::thnn_conv2d 0.000us 10.484ms191# aten::thnn_conv2d_forward 10.484ms 10.484ms192# void at::native::im2col_kernel<float>(long, float co... 3.844ms 3.844ms193# sgemm_32x32x32_NN 3.206ms 3.206ms194# sgemm_32x32x32_NN_vec 3.093ms 3.093ms195# ------------------------------------------------------- ------------ ------------196# Self CPU time total: 23.015ms197# Self CUDA time total: 11.666ms198#199200######################################################################201# Note the occurrence of on-device kernels in the output (e.g. ``sgemm_32x32x32_NN``).202203######################################################################204# 4. Using profiler to analyze memory consumption205# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~206#207# PyTorch profiler can also show the amount of memory (used by the model's tensors)208# that was allocated (or released) during the execution of the model's operators.209# In the output below, 'self' memory corresponds to the memory allocated (released)210# by the operator, excluding the children calls to the other operators.211# To enable memory profiling functionality pass ``profile_memory=True``.212213model = models.resnet18()214inputs = torch.randn(5, 3, 224, 224)215216with profile(activities=[ProfilerActivity.CPU],217profile_memory=True, record_shapes=True) as prof:218model(inputs)219220print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))221222# (omitting some columns)223# --------------------------------- ------------ ------------ ------------224# Name CPU Mem Self CPU Mem # of Calls225# --------------------------------- ------------ ------------ ------------226# aten::empty 94.79 Mb 94.79 Mb 121227# aten::max_pool2d_with_indices 11.48 Mb 11.48 Mb 1228# aten::addmm 19.53 Kb 19.53 Kb 1229# aten::empty_strided 572 b 572 b 25230# aten::resize_ 240 b 240 b 6231# aten::abs 480 b 240 b 4232# aten::add 160 b 160 b 20233# aten::masked_select 120 b 112 b 1234# aten::ne 122 b 53 b 6235# aten::eq 60 b 30 b 2236# --------------------------------- ------------ ------------ ------------237# Self CPU time total: 53.064ms238239print(prof.key_averages().table(sort_by="cpu_memory_usage", row_limit=10))240241#############################################################################242# The output might look like this (omitting some columns):243#244# .. code-block:: sh245#246# --------------------------------- ------------ ------------ ------------247# Name CPU Mem Self CPU Mem # of Calls248# --------------------------------- ------------ ------------ ------------249# aten::empty 94.79 Mb 94.79 Mb 121250# aten::batch_norm 47.41 Mb 0 b 20251# aten::_batch_norm_impl_index 47.41 Mb 0 b 20252# aten::native_batch_norm 47.41 Mb 0 b 20253# aten::conv2d 47.37 Mb 0 b 20254# aten::convolution 47.37 Mb 0 b 20255# aten::_convolution 47.37 Mb 0 b 20256# aten::mkldnn_convolution 47.37 Mb 0 b 20257# aten::max_pool2d 11.48 Mb 0 b 1258# aten::max_pool2d_with_indices 11.48 Mb 11.48 Mb 1259# --------------------------------- ------------ ------------ ------------260# Self CPU time total: 53.064ms261#262263######################################################################264# 5. Using tracing functionality265# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~266#267# Profiling results can be outputted as a ``.json`` trace file:268269model = models.resnet18().cuda()270inputs = torch.randn(5, 3, 224, 224).cuda()271272with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:273model(inputs)274275prof.export_chrome_trace("trace.json")276277######################################################################278# You can examine the sequence of profiled operators and CUDA kernels279# in Chrome trace viewer (``chrome://tracing``):280#281# .. image:: ../../_static/img/trace_img.png282# :scale: 25 %283284######################################################################285# 6. Examining stack traces286# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~287#288# Profiler can be used to analyze Python and TorchScript stack traces:289290with profile(291activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],292with_stack=True,293) as prof:294model(inputs)295296# Print aggregated stats297print(prof.key_averages(group_by_stack_n=5).table(sort_by="self_cuda_time_total", row_limit=2))298299#################################################################################300# The output might look like this (omitting some columns):301#302# .. code-block:: sh303#304# ------------------------- -----------------------------------------------------------305# Name Source Location306# ------------------------- -----------------------------------------------------------307# aten::thnn_conv2d_forward .../torch/nn/modules/conv.py(439): _conv_forward308# .../torch/nn/modules/conv.py(443): forward309# .../torch/nn/modules/module.py(1051): _call_impl310# .../site-packages/torchvision/models/resnet.py(63): forward311# .../torch/nn/modules/module.py(1051): _call_impl312# aten::thnn_conv2d_forward .../torch/nn/modules/conv.py(439): _conv_forward313# .../torch/nn/modules/conv.py(443): forward314# .../torch/nn/modules/module.py(1051): _call_impl315# .../site-packages/torchvision/models/resnet.py(59): forward316# .../torch/nn/modules/module.py(1051): _call_impl317# ------------------------- -----------------------------------------------------------318# Self CPU time total: 34.016ms319# Self CUDA time total: 11.659ms320#321322######################################################################323# Note the two convolutions and the two call sites in ``torchvision/models/resnet.py`` script.324#325# (Warning: stack tracing adds an extra profiling overhead.)326327######################################################################328# 7. Using profiler to analyze long-running jobs329# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~330#331# PyTorch profiler offers an additional API to handle long-running jobs332# (such as training loops). Tracing all of the execution can be333# slow and result in very large trace files. To avoid this, use optional334# arguments:335#336# - ``schedule`` - specifies a function that takes an integer argument (step number)337# as an input and returns an action for the profiler, the best way to use this parameter338# is to use ``torch.profiler.schedule`` helper function that can generate a schedule for you;339# - ``on_trace_ready`` - specifies a function that takes a reference to the profiler as340# an input and is called by the profiler each time the new trace is ready.341#342# To illustrate how the API works, let's first consider the following example with343# ``torch.profiler.schedule`` helper function:344345from torch.profiler import schedule346347my_schedule = schedule(348skip_first=10,349wait=5,350warmup=1,351active=3,352repeat=2)353354######################################################################355# Profiler assumes that the long-running job is composed of steps, numbered356# starting from zero. The example above defines the following sequence of actions357# for the profiler:358#359# 1. Parameter ``skip_first`` tells profiler that it should ignore the first 10 steps360# (default value of ``skip_first`` is zero);361# 2. After the first ``skip_first`` steps, profiler starts executing profiler cycles;362# 3. Each cycle consists of three phases:363#364# - idling (``wait=5`` steps), during this phase profiler is not active;365# - warming up (``warmup=1`` steps), during this phase profiler starts tracing, but366# the results are discarded; this phase is used to discard the samples obtained by367# the profiler at the beginning of the trace since they are usually skewed by an extra368# overhead;369# - active tracing (``active=3`` steps), during this phase profiler traces and records data;370# 4. An optional ``repeat`` parameter specifies an upper bound on the number of cycles.371# By default (zero value), profiler will execute cycles as long as the job runs.372373######################################################################374# Thus, in the example above, profiler will skip the first 15 steps, spend the next step on the warm up,375# actively record the next 3 steps, skip another 5 steps, spend the next step on the warm up, actively376# record another 3 steps. Since the ``repeat=2`` parameter value is specified, the profiler will stop377# the recording after the first two cycles.378#379# At the end of each cycle profiler calls the specified ``on_trace_ready`` function and passes itself as380# an argument. This function is used to process the new trace - either by obtaining the table output or381# by saving the output on disk as a trace file.382#383# To send the signal to the profiler that the next step has started, call ``prof.step()`` function.384# The current profiler step is stored in ``prof.step_num``.385#386# The following example shows how to use all of the concepts above:387388def trace_handler(p):389output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)390print(output)391p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")392393with profile(394activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],395schedule=torch.profiler.schedule(396wait=1,397warmup=1,398active=2),399on_trace_ready=trace_handler400) as p:401for idx in range(8):402model(inputs)403p.step()404405406######################################################################407# Learn More408# ----------409#410# Take a look at the following recipes/tutorials to continue your learning:411#412# - `PyTorch Benchmark <https://pytorch.org/tutorials/recipes/recipes/benchmark.html>`_413# - `PyTorch Profiler with TensorBoard <https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html>`_ tutorial414# - `Visualizing models, data, and training with TensorBoard <https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html>`_ tutorial415#416417418