XProf Profiler Plugin

Warnings

• gotcha A regression in `libtpu` versions `0.0.35` and `0.0.37` affects HLO Module-dependent tools like 'HLO Op Profile', 'Trace Viewer', and 'Graph Viewer' in XProf. It is recommended to use `libtpu 0.0.36` as a temporary workaround.
  Fix: Downgrade or upgrade `libtpu` to `0.0.36` or wait for a fix in newer `libtpu` versions.
• gotcha The TensorBoard Profiler Plugin requires internet access to load the Google Chart library. If running TensorBoard offline, behind a corporate firewall, or in a datacenter, some charts and tables in the profiler interface may not display correctly.
  Fix: Ensure internet connectivity or configure firewall rules to allow access to Google Chart libraries. For corporate environments, consider using the `--bind_all` TensorBoard flag if applicable.
• gotcha When using virtual environments, ensure that both `TensorBoard` and `tensorboard-plugin-profile` are installed within the *same and activated* virtual environment. Mixing installations or not activating the environment can lead to the profiler tab not appearing or displaying a 'plugin has moved' error.
  Fix: Always install `tensorboard` and `tensorboard-plugin-profile` in the same virtual environment and ensure it's activated before launching TensorBoard.
• breaking The Profiler plugin requires recent versions of TensorFlow and TensorBoard. Specifically, `TensorFlow >= 2.18.0` and `TensorBoard >= 2.18.0` are prerequisites. Older versions may lead to compatibility issues, including 'The profile plugin has moved' messages even after installation.
  Fix: Upgrade TensorFlow and TensorBoard to versions `2.18.0` or newer: `pip install --upgrade tensorflow tensorboard`.
• gotcha For GPU profiling, the NVIDIA CUDA Profiling Tools Interface (CUPTI) must be correctly configured and accessible via the `LD_LIBRARY_PATH` environment variable. Insufficient privileges or an incorrect path can prevent GPU profiling data collection.
  Fix: Verify CUPTI installation and path using `/sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti`. If not found, prepend its installation directory (e.g., `/usr/local/cuda/extras/CUPTI/lib64`) to `LD_LIBRARY_PATH`.
• gotcha Running the profiler for excessively long durations can lead to out-of-memory errors. It is recommended to profile no more than 10 steps at a time. Also, avoid profiling the first few batches of training, as initialization overhead can skew results.
  Fix: Limit profiling duration or step range (e.g., `profile_batch='10,15'` in Keras callback) and skip initial warm-up steps.

Install

• pip install tensorboard-plugin-profile Install latest stable version

Quickstart

This quickstart demonstrates how to generate profiling data for a TensorFlow model and then launch TensorBoard to visualize it. The `tensorboard-plugin-profile` automatically integrates with TensorBoard once installed. After running the Python code, open a new terminal and execute the `tensorboard` command provided to view the profile data in your browser under the 'Profile' tab.

import tensorflow as tf
from datetime import datetime
import os

# Ensure log directory exists
log_dir = os.path.join("logs", "profile", datetime.now().strftime("%Y%m%d-%H%M%S"))
os.makedirs(log_dir, exist_ok=True)

# Dummy model and data for profiling
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy')
data = tf.random.normal(shape=(100, 10))
labels = tf.random.uniform(shape=(100, 1), maxval=2, dtype=tf.int64)

# Option 1: Programmatic profiling with tf.profiler.experimental
print(f"Starting programmatic profile, data will be in {log_dir}")
with tf.profiler.experimental.Profile(log_dir):
    model.fit(data, labels, epochs=2, batch_size=32)

# Option 2: Using TensorBoard Keras Callback for profiling specific batches
# tb_callback = tf.keras.callbacks.TensorBoard(
#     log_dir=log_dir,
#     profile_batch='1,3' # Profile batches 1 to 3
# )
# model.fit(data, labels, epochs=2, batch_size=32, callbacks=[tb_callback])

print("Profiling data generated. To view, run TensorBoard in your terminal:")
print(f"tensorboard --logdir={os.path.abspath('logs')}")
print("Then open your browser to http://localhost:6006/#profile")