Profilers

Core Module

Profilers

In general profiling code is about improving the performance of your code. In this session we are going to take a somewhat narrow approach to what "performance" is: runtime, meaning the time it takes to execute your program.

At the bare minimum, the two questions a proper profiling of your program should be able to answer is:

“ How many times is each method in my code called?”
“ How long do each of these methods take?”

The first question is important to priorities optimization. If two methods A and B have approximately the same runtime, but A is called 1000 more times than B we should probably spend time optimizing A over B if we want to speedup our code. The second question is gives itself, directly telling us which methods are the expensive to call.

Using profilers can help you find bottlenecks in your code. In this exercise we will look at two different profilers, with the first one being the cProfile. cProfile is pythons build in profiler that can help give you an overview runtime of all the functions and methods involved in your programs.

❔ Exercises

Run the cProfile on the vae_mnist_working.py script. Hint: you can directly call the profiler on a script using the -m arg
```
python -m cProfile -o <output_file> -s <sort_order> myscript.py
```
Try looking at the output of the profiling. Can you figure out which function took the longest to run?
Can you explain the difference between tottime and cumtime? Under what circumstances does these differ and when are they equal.
To get a better feeling of the profiled result we can try to visualize it. Python does not provide a native solution, but open-source solutions such as snakeviz exist. Try installing snakeviz and load a profiled run into it (HINT: snakeviz expect the run to have the file format .prof).
Try optimizing the run! (Hint: The data is not stored as torch tensor). After optimizing the code make sure (using cProfile and snakeviz) that the code actually runs faster.

PyTorch profiling

Profiling machine learning code can become much more complex because we are suddenly beginning to mix different devices (CPU+GPU), that can (and should) overlap some of their computations. When profiling this kind of machine learning code we are often looking for bottlenecks. A bottleneck is simple the place in your code that is preventing other processes from performing their best. This is the reason that all major deep learning frameworks also include their own profilers that can help profiling more complex applications.

The image below show a typical report using the build in profiler in pytorch. As the image shows the profiler looks both a the kernel time (this is the time spend doing actual computations) and also transfer times such as memcpy (where we are copying data between devices). It can even analyze your code and give recommendations.

Using the profiler can be as simple as wrapping the code that you want to profile with the torch.profiler.profile decorator

with torch.profiler.profile(...) as prof:
    # code that I want to profile
    output = model(data)

❔ Exercises

Exercise files

In these investigate the profiler that is build into PyTorch already. Note that these exercises requires that you have PyTorch v1.8.1 installed (or higher). You can always check which version you currently have installed by writing (in a python interpreter):

import torch
print(torch.__version__)

But we always recommend to update to the latest PyTorch version for the best experience. Additionally, to display the result nicely (like snakeviz for cProfile) we are also going to use the tensorboard profiler extension

pip install torch_tb_profiler

A good starting point is too look at the API for the profiler. Here the important class to look at is the torch.profiler.profile class.

Lets try out an simple example (taken from here):

Try to run the following code

import torch
import torchvision.models as models
from torch.profiler import profile, ProfilerActivity

model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    model(inputs)

this will profile the forward pass of Resnet 18 model.

Running this code will produce an prof object that contains all the relevant information about the profiling. Try writing the following code:
```
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
```
what operation is taking most of the cpu?

Try running

print(prof.key_averages(group_by_input_shape=True).table(sort_by="cpu_time_total", row_limit=30))

can you see any correlation between the shape of the input and the cost of the operation?

(Optional) If you have a GPU you can also profile the operations on that device:

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    model(inputs)

(Optional) As an alternative to using profile as an context-manager we can also use its .start and .stop methods:
```
prof = profile(...)
prof.start()
...  # code I want to profile
prof.stop()
```
Try doing this on the above example.

The torch.profiler.profile function takes some additional arguments. What argument would you need to set to also profile the memory usage? (Hint: this page) Try doing it to the simple example above and make sure to sort the sample by self_cpu_memory_usage.
As mentioned we can also get a graphical output for better inspection. After having done a profiling try to export the results with:
```
prof.export_chrome_trace("trace.json")
```
you should be able to visualize the file by going to chrome://tracing in any chromium based web browser. Can you still identify the information printed in the previous exercises from the visualizations?
Running profiling on a single forward step can produce misleading results as it only provides a single sample that may depend on what background processes that are running on your computer. Therefore it is recommended to profile multiple iterations of your model. If this is the case then we need to include prof.step() to tell the profiler when we are doing a new iteration
```
with profile(...) as prof:
    for i in range(10):
        model(inputs)
        prof.step()
```
Try doing this. Is the conclusion this the same on what operations that are taken up most of the time? Have the percentage changed significantly?
Additionally, we can also visualize the profiling results using the profiling viewer in tensorboard.
1. Start by initializing the profile class with an additional argument:
```
from torch.profiler import profile, tensorboard_trace_handler
with profile(..., on_trace_ready=tensorboard_trace_handler("./log/resnet18")) as prof:
    ...
```
  Try run a profiling (using a couple of iterations) and make sure that a file with the .pt.trace.json is produced in the log/resnet18 folder.
2. Now try launching tensorboard
```
tensorboard --logdir=./log
```
  and open the page http://localhost:6006/#pytorch_profiler, where you should hopefully see an image similar to the one below:
  
  Image credit
  
  Try poking around in the interface.
3. Tensorboard have a nice feature for comparing runs under the diff tab. Try redoing a profiling run but use model = models.resnet34() instead. Load up both runs and try to look at the diff between them.
As an final exercise, try to use the profiler on the vae_mnist_working.py file from the previous module on debugging, where you profile a whole training run (not only the forward pass). What is the bottleneck during the training? Is it still the forward pass or is it something else? Can you improve the code somehow based on the information from the profiler.

This end the module on profiling. If you want to go into more details on this topic we can recommend looking into line_profiler and kernprof. A downside of using python's cProfile is that it can only profiling at an functional/modular level, that is great for identifying hotspots in your code. However, sometimes the cause of an computationally hotspot is a single line of code in a function, which will not be caught by cProfile. An example would be an simple index operations such as a[idx] = b, which for large arrays and non-sequential indexes is really expensive. For these cases line_profiler and kernprof are excellent tools to have in your toolbox. Additionally, if you do not like cProfile we can also recommend py-spy which is another open-source profiling tool for Python programs.