HowTo profile TensorFlow:. Nowadays TensorFlow one of the most... | by Illarion Khlestov | Towards Data Science

Nowadays TensorFlow one of the most used library for machine learning. Sometimes it may be quite useful to profile tensorflow graph and know what operations take more time and what less. This can be done with tensorflow timeline module. Unfortunately, I cannot find any clear tutorial how to use it. So in this blog post, I will try to resolve this issue with such topics covered:
How to perform profiling of tensorflow code.
How to merge timelines from multiply session runs.
What issues may occur during profiling and how to resolve them.
First let’s define simple example, followed by StackOverflow answer:
You should note additional options and run_metadata provided to the session run. This script should work as on CPU, as on GPU. After execution, we will have a timeline_01.json file with our profiled data stored in Chrome trace format. If your script failed — try the first solution from Issues during profiling section.
To view stored data, we should use Chrome browser(unfortunately only it supports own tracing format, as far as I know). Go to the page chrome://tracing. In the upper left corner, you will find Load button. Press it and load our JSON file.
On the top you will see time axis in ms. To get more precise info about some operation just click on it. Also on the right side, there are simple tools exist: selection, pan, zoom and timing.
Now let’s define more complicated example with some placeholders and optimizer:
Now our operations stored under the variable scopes. With such approach, operations names will be started with the scope name and clearly distinguished on the timeline.
Also, code store traces for three runs. If we execute the script on CPU we receive three relatively similar timelines such this:
But if we check results from GPU profiling first results will be different from next ones:
You may note that first run takes on GPU much more time than later ones. This happens because on first run tensorflow performs some GPU initialization routines, and later they will be optimized. If you want the more accurate timeline, you should store the tracings after one hundred runs or so.
Also now all incoming/outcoming flows are started with variable scope name, and we know exactly where one or another operation exist in the source code.
What if we want to store multiply session runs in one file for some reason? Unfortunately, this can be done only manually. Inside Chrome trace format there are stored definitions for every event and it’s running time. At the first iteration we will store all data, but on next runs, we will only update running time, not definitions itself. Here is only the class definition for merging events, and the full example you may find here:
And we receive such cool merged timeline:
It seems that initialization still takes a lot, so let’s zoom to the right side:
Now we can see some repeated pattern. There no any particular separators between runs, but we may visually distinguish them.
There are some troubles may exist during profiling. First of all, it may not work at all. If you’ve met some error like this:
I tensorflow/stream_executor/dso_loader.cc:126] Couldn't open CUDA library libcupti.so.8.0. LD_LIBRARY_PATH:
and you sure that without profiling all works as expected it can be resolved according to the GitHub issue with the installation of the additional library libcupti-dev. This command should fix described error:
sudo apt-get install libcupti-dev
Second is delays during runs. On the last image we saw that there is a gap between runs. For large networks, this can take really large amount of time. This bug cannot be fully resolved, but the delays may be decreased with a custom C++ protobuf library. It’s clearly described in the tensorflow docs how to perform an installation.
I hope with such profiling you will understand more deeply what’s going in tensorflow under the hood and what parts of your graph may be optimized. All code examples with already generated timelines on CPU and GPU are stored in this repo.
Thanks for reading!