How to profile TensorFlow training in TensorBoard

TensorFlow training profiles show where batches spend time across input pipeline work, model execution, and host overhead. TensorBoard reads profiler traces from the training log directory, so the trace can point investigation toward slow input stages, long model kernels, or host-side gaps between batches.

The explicit tf.profiler.experimental.start() and tf.profiler.experimental.stop() calls capture the section of training that needs inspection. That boundary avoids relying on Keras callback profile capture, whose profile_batch behavior has changed across TensorFlow and Keras releases.

Run TensorBoard from the same Python environment that owns the training log directory. The profile plugin must be installed for the Profile tab, and minimal Python environments may need a setuptools constraint because TensorBoard releases that still import pkg_resources fail after that module is removed.

1. Install TensorBoard and the profile plugin in the active TensorFlow environment.

   $ python3 -m pip install --upgrade tensorboard tensorboard-plugin-profile "setuptools<81"

   The setuptools<81 constraint keeps the pkg_resources module available for TensorBoard releases that still import it. Remove the constraint only after the installed TensorBoard version starts without that module.

   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Save a small training profile script as

   profile-training.py

   .

   profile-training.py

   import os
   import shutil
   from pathlib import Path
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
    
   logdir = Path("logs/profile-demo")
   if logdir.exists():
       shutil.rmtree(logdir)
   logdir.mkdir(parents=True)
    
   features = tf.random.uniform((512, 10), seed=42)
   labels = tf.cast(tf.reduce_sum(features, axis=1, keepdims=True) > 5.0, tf.float32)
   dataset = (
       tf.data.Dataset.from_tensor_slices((features, labels))
       .shuffle(512, seed=42)
       .batch(32)
       .prefetch(tf.data.AUTOTUNE)
   )
    
   model = tf.keras.Sequential(
       [
           tf.keras.layers.Input(shape=(10,)),
           tf.keras.layers.Dense(16, activation="relu"),
           tf.keras.layers.Dense(1, activation="sigmoid"),
       ]
   )
   model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    
   tf.profiler.experimental.start(str(logdir))
   try:
       model.fit(dataset, epochs=2, verbose=0, shuffle=False)
   finally:
       tf.profiler.experimental.stop()
    
   profile_files = sorted(logdir.glob("plugins/profile/*/*"))
   profile_dirs = sorted({profile_file.parent for profile_file in profile_files})
    
   print(f"TensorBoard log directory: {logdir}")
   print(f"Profiler files: {len(profile_files)}")
   for profile_dir in profile_dirs:
       print(f"Profile directory: {profile_dir}")
   for profile_file in profile_files:
       print(f"Trace file: {profile_file.name}")

   The script removes only its own logs/profile-demo directory before each run so stale profile files cannot hide a failed capture.

   Related: How to run a custom training loop in TensorFlow

3. Run the training script and confirm that one profiler trace is written.

   $ python3 profile-training.py
   TensorBoard log directory: logs/profile-demo
   Profiler files: 1
   Profile directory: logs/profile-demo/plugins/profile/2026_06_29_04_29_06
   Trace file: training-worker.xplane.pb

   The profile directory timestamp and xplane.pb filename are generated for each run, so the exact values differ between machines.

4. Start TensorBoard against the profile log directory.

   $ tensorboard --logdir logs/profile-demo --host 127.0.0.1 --port 6006
   TensorBoard 2.20.0 at http://127.0.0.1:6006/ (Press CTRL+C to quit)

   Keep --host 127.0.0.1 for local analysis. Binding TensorBoard to 0.0.0.0 exposes the dashboard to other hosts that can reach the machine.

5. Open the local Profile page in a browser.

   http://127.0.0.1:6006/#profile

   The Profile tab should be available and the run should contain the profile-demo trace. Use Overview Page, Input Pipeline Analyzer, and Trace Viewer to compare input time, host time, and model execution.