TensorFlow Profiler offers guidance for faster model training [code included]

TensorFlow 2.0 profiler can help improving your model performance.

TensorFlow Profiler provides a set of tools that you can use to measure the training performance and resource consumption of your TensorFlow models.

Beyond the ability to capture and investigate numerous aspects of a profile, the tools offer guidance on how to resolve performance bottlenecks

New profiling tools:

1. Overview Page: Provides a top-level view of model performance and recommendations to optimize performance
2. Input Pipeline Analyzer: Analyzes your model’s data input pipeline for bottlenecks and recommends improvements to improve performance
3. TensorFlow Stats: Displays performance statistics for every TensorFlow operation executed during the profiling session
4. GPU Kernel Stats: Displays performance statistics and the originating operation for every GPU accelerated kernel

1- Install the Profiler plugin for TensorBoard:

pip install -U tensorboard_plugin_profile

This adds the full Profiler capabilities to our TensorBoard installation.

2- TensorBoard callback:

Next, enable our model training to capture a profile using the TensorBoard callback in Keras:

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir = logs,
                                                      profile_batch = '500,510')

We can choose which batches to profile with the profile_batch parameter.

We should skip the first few batches to avoid inaccuracy due to initialization overhead. 

3- Start TensorBoard:

tensorboard --logdir {log directory}    # in terminal
%tensorboard --logdir {log directory}   # in Colab

Example:

In this example, we explore the capabilities of the TensorFlow Profiler by capturing the performance profile obtained by training a model to classify images in the MNIST dataset.

4- Recommendations

Not only can the Profiler tell you where your program has bottlenecks, it can tell you what you can do to resolve them and make your code execute faster.

5- Applying the recommendations:

5- Results

As shown in the figure below, we have a much better speed.

By following the instructions for some toy tutorials, we can achieve up to 10x speed improvement! 

References

• Profiler tutorial in Colab
• Guide
• Video TensorFlow Dev Summit 2020
• Demo

Best Regards