Optimizing Model Training with TensorFlow Profiler | by Rohan Jagtap | Towards Data Science

We want our models to train real fast. We use GPUs to make operations execute faster. However, it is possible that even after speeding up the computations, the model may have inefficiencies in the pipeline itself, and thus, may train slower. In such cases, it becomes really difficult to debug the code, or as a matter of fact, even tell what is slow.
This can be addressed by using the TensorFlow Profiler. The Profiler ‘profiles’ the TensorFlow code execution. We’ll be discussing the Profiler, how to use it, best practices, and how to optimize the GPU performance in this article.
Note that this article expects basic knowledge of training TensorFlow models and using Tensorboard. You can refer to my article on Tensorboard if you’re unaware of it.
First things first, even before optimizing anything, let’s talk about what the profiler is and what it offers.
Profiling helps you understand the hardware resource consumption (time and memory) of the various TensorFlow operations (ops) in your model and resolve performance bottlenecks and ultimately, make the model execute faster.
— TensorFlow Docs
Basically, the profiler monitors the model training. It notes down the time required for the ops to execute, the time required for complete steps to execute, collects insights on resource utilization in terms of time and memory and provides visualizations for understanding these insights.
In the upcoming sections, we’ll see how to work with the profiler:
We’ll train a simple model with profiling.Explore the profiler to gain insights into the training.Use these insights to optimize training.
We’ll train a simple model with profiling.
Explore the profiler to gain insights into the training.
Use these insights to optimize training.
To understand how to use the profiler, let’s first train a model! We’ll use the famous mnist dataset from tensorflow-datasets:
To install TensorFlow Profiler, use:
pip install -U tensorboard_plugin_profile
Note that Profiler requires the latest versions of TensorFlow and TensorBoard installed.
Now, load and preprocess the dataset:
(ds_train, ds_test), ds_info = tfds.load(      'mnist',      split=['train', 'test'],      shuffle_files=True,as_supervised=True,with_info=True,                       )def rescale(image, label):      return tf.cast(image, tf.float32) / 255., labelds_train = ds_train.map(rescale)ds_train = ds_train.batch(128)
Build the model:
model = tf.keras.models.Sequential([      tf.keras.layers.Flatten(input_shape=(28, 28, 1))      tf.keras.layers.Dense(128,activation='relu'),      tf.keras.layers.Dense(10, activation='softmax')])model.compile(      loss='sparse_categorical_crossentropy',      optimizer=tf.keras.optimizers.Adam(0.001),      metrics=['accuracy'])
Now, enabling the profiler is super easy. You simply need to add an extra argument to the tensorboard callback:
tensorboard_callback = tf.keras.callbacks.TensorBoard(      log_dir = logs,      histogram_freq = 1,      profile_batch = '500,520')
The profile_batch argument tells TensorFlow to capture the execution performance and profile it. ‘500,520’ specifies the batches you wish to profile, i.e., the model’s performance from batch 500 through 520 will be captured.
Now, we’ll train the model as usual with the above callback:
model.fit(      ds_train,      epochs=2,      validation_data=ds_test,      callbacks = [tensorboard_callback])
Now that we have trained the model, let’s see how we did with the help of the profiler:
This section gives a high-level summary of the model performance. Here are a few things to look for:
The Step-time Graph (right) is a plot of the step number (x-axis) v/s the time required to execute the corresponding step (step-time, y-axis). It also shows what portion of this step-time is utilized by each category (the colored partitions). Looking at the above graph, we can tell that our model is ‘HIGHLY input bound,’ i.e., it spends a lot of its training time in input operations. We’ll optimize this in some time:)The Average Step Time (under Performance Summary) gives a break-down of the average step-time. In an ideal case, we want our model to spend most of its time in actually ‘training,’ i.e., keeping the GPU busy (Device Compute Time must be the highest, and all the other overheads must be as low as possible).TF Op Placement (under Performance Summary) shows the % of operations being executed on the host (CPU) v/s the device (GPU). To maximize GPU utilization, the device Op placement must be maximum.
The Step-time Graph (right) is a plot of the step number (x-axis) v/s the time required to execute the corresponding step (step-time, y-axis). It also shows what portion of this step-time is utilized by each category (the colored partitions). Looking at the above graph, we can tell that our model is ‘HIGHLY input bound,’ i.e., it spends a lot of its training time in input operations. We’ll optimize this in some time:)
The Average Step Time (under Performance Summary) gives a break-down of the average step-time. In an ideal case, we want our model to spend most of its time in actually ‘training,’ i.e., keeping the GPU busy (Device Compute Time must be the highest, and all the other overheads must be as low as possible).
TF Op Placement (under Performance Summary) shows the % of operations being executed on the host (CPU) v/s the device (GPU). To maximize GPU utilization, the device Op placement must be maximum.
In an ideal case, your program should have high GPU utilization, minimal CPU (host) to GPU (device) communication, and no overhead from the input pipeline.
— TensorFlow Docs
So before this tool gets anymore intimidating, let’s break it down:
to the left (vertical grey column), we can see two major sections: /device and /host. This tells us which TensorFlow op executed on which device (GPU or CPU resp.).
to the right, the colored bars are the durations for which the respective TensorFlow ops have executed.
Typically, the host executes input operations, preprocesses training data and transfers it to the device, while the device executes the actual model training.
— TensorFlow Docs
Exploring this tool will give you a better idea of these durations (Use W and S to zoom in or out, respectively). For starters, take a look at this:
Stream #19 shows the timeline used to launch compute kernels and make device to device copies.Stream #20 is used for host to device copies andStream#21 for device to host copies.
Stream #19 shows the timeline used to launch compute kernels and make device to device copies.
Stream #20 is used for host to device copies and
Stream#21 for device to host copies.
Ideally, the GPU compute timeline must be crowded (i.e., GPU busy). Host to device and device to host copies should be minimal.
Using the Steps timeline, you can see the different phases of the model like the forward pass (sequential and dense), loss calculation (sparse_categorical_crossentropy), and backward pass (gradient_tape and Adam). These ops should be crowded with minimal space between steps. The spaces indicate GPU idle time, and crowding indicates utilization.
The rest of the tools give comprehensive insights on specific categories or ops covered in the Overview Page and Trace Viewer. Let’s quickly glance through them:
Input Pipeline Analyzer: The input pipeline analyzer checks the input pipeline and informs whether there is a performance bottleneck. It also tells us if our model is input bound.GPU Kernel Stats: This tool shows performance statistics and the originating op for every GPU kernel.Memory Profile: This tool profiles the GPU memory usage. This can be used to analyze and debug the OOM (Out of Memory) error, raised when the GPU’s memory is exhausted.TensorFlow Stats: This tool gives a performance overview of every TensorFlow op that is executed. It provides a decent grouping for the host and device ops for readability.
Input Pipeline Analyzer: The input pipeline analyzer checks the input pipeline and informs whether there is a performance bottleneck. It also tells us if our model is input bound.
GPU Kernel Stats: This tool shows performance statistics and the originating op for every GPU kernel.
Memory Profile: This tool profiles the GPU memory usage. This can be used to analyze and debug the OOM (Out of Memory) error, raised when the GPU’s memory is exhausted.
TensorFlow Stats: This tool gives a performance overview of every TensorFlow op that is executed. It provides a decent grouping for the host and device ops for readability.
In the previous sections, we trained our model and saw how the profiler helps in identifying potential inefficiencies in the pipeline and resource utilization. This section will try and address those inefficiencies and see if our model does any better...
Now, we’ve seen that our model is highly ‘input bound,’ i.e., most of its training time was spent in input operations (e.g., preprocessing). If you observe tf_data_iterator_get_next from the trace viewer, we can see that tf.data operations account for a lot of time. This is because our tf.data.Dataset fetches the batches in real-time during training.
One thing we can do to optimize this is use caching and prefetching.
tf.data.Dataset.cache will load the dataset in memory or in local storage. This will save time for expensive I/O operations like opening and reading the data from files on each epoch.tf.data.Dataset.prefetch will essentially prepare the data in advance and in parallel.
tf.data.Dataset.cache will load the dataset in memory or in local storage. This will save time for expensive I/O operations like opening and reading the data from files on each epoch.
tf.data.Dataset.prefetch will essentially prepare the data in advance and in parallel.
So without prefetch, at a given step s, first the data is fetched (and preprocessed) in time t1. Then, the model is trained on this data in time t2. Now the total step time for step s becomes (t1 + t2).
But with prefetch, a background thread fetches the data for step (s + 1) while the model is being trained on the data for step s. Hence, if the time required to fetch the data for step (s + 1) is t1, and the time required for training the model for step s is t2, then the step time will be max(t1, t2).
Let’s add caching and prefetching to our previous code:
ds_train = ds_train.map(rescale)ds_train = ds_train.batch(128)ds_train = ds_train.cache()ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)
And then, train the model in the exact same way like before.
Now that we have optimized our training, let’s take a look at the profiler:
It can be clearly seen that the input overhead has significantly decreased. Moreover, the average step time is almost halved (6.2 to 3.2). You can see in the recommendations that the program is ‘MODERATELY input bound.’
The device TF Op placement has also fairly increased.
Let’s check trace viewer:
Notice that Stream #19 is more crowded now, which means GPU utilization has improved.
Stream #20 and #21 have been minimized, so the host to device and device to host copies have decreased.
tf_data_iterator_get_next is also less crowded this time since we have used caching and prefetching.
In the previous section, we saw how optimization can result in better model performance. There I made just one optimization change for the sake of demonstration. But there’s a lot more to optimization than that.
In this section, we’ll glance through some optimizations to look out for:
First, USING A GPU ISN’T ALWAYS GOOD. In the previous sections, we’ve seen that offloading computations to a GPU accounts for the device to host and host to device copying overheads. Moreover, launching GPU kernels also adds latency. So, if your model is so small that the time required for these overheads is more than the actual training time, is using the GPU really the right thing?In the trace viewer, we saw that the gaps between steps should be small (gaps indicate GPU idle time). We optimized the input pipeline in the previous section, but we couldn’t see much difference in the gap lengths. One of the possible reasons for this is: the GPU does some host-side activities like copying data or scheduling at the beginning of each step. At the same time, tf.data is operating on threads to parallelize the input pipeline at the CPU-side. There is a chance that these CPU-side threads may interfere with the GPU activity. For this, you can set the environment variable: TF_GPU_THREAD_MODE=gpu_private. This ensures that GPU kernels are launched from their own dedicated threads and don’t get queued behind tf.data work.Even after debugging threads, if the gaps between the steps do not shrink, check your callbacks, executions after steps, and metrics evaluation. If disabling these improve performance, then you might want to consider computing these only on a few steps. Moreover, for custom training, you can use tf.while_loop .We’ve already discussed that ideally, most of the ops must be placed on the GPU. TF Op Placement on the Overview page helps us with that. This reduces the host-device communication for making copies. Here’s a trick for ensuring that the ops are optimally placed on the GPU: even for one GPU, use a distribution strategy like tf.distribute.OneDeviceStrategy .
First, USING A GPU ISN’T ALWAYS GOOD. In the previous sections, we’ve seen that offloading computations to a GPU accounts for the device to host and host to device copying overheads. Moreover, launching GPU kernels also adds latency. So, if your model is so small that the time required for these overheads is more than the actual training time, is using the GPU really the right thing?
In the trace viewer, we saw that the gaps between steps should be small (gaps indicate GPU idle time). We optimized the input pipeline in the previous section, but we couldn’t see much difference in the gap lengths. One of the possible reasons for this is: the GPU does some host-side activities like copying data or scheduling at the beginning of each step. At the same time, tf.data is operating on threads to parallelize the input pipeline at the CPU-side. There is a chance that these CPU-side threads may interfere with the GPU activity. For this, you can set the environment variable: TF_GPU_THREAD_MODE=gpu_private. This ensures that GPU kernels are launched from their own dedicated threads and don’t get queued behind tf.data work.
Even after debugging threads, if the gaps between the steps do not shrink, check your callbacks, executions after steps, and metrics evaluation. If disabling these improve performance, then you might want to consider computing these only on a few steps. Moreover, for custom training, you can use tf.while_loop .
We’ve already discussed that ideally, most of the ops must be placed on the GPU. TF Op Placement on the Overview page helps us with that. This reduces the host-device communication for making copies. Here’s a trick for ensuring that the ops are optimally placed on the GPU: even for one GPU, use a distribution strategy like tf.distribute.OneDeviceStrategy .
In this rather very exhaustive article, we discussed a very useful TensorFlow tool, which can be used to debug a model. We saw what the tool offers, the basics of using it, and finally, ways to optimize a TensorFlow model’s performance on GPU.
There’s a lot more to the profiler and optimization, and I encourage you to read the docs (link in references) for more!
One important optimization technique that I didn’t touch in this article was using XLA and fp16 Mixed Precision. This can actually take a while to cover, and best be explained in a separate article!