PyTorch Profiler TensorBoard Plugin

Warnings

• breaking The TensorBoard integration with PyTorch profiler (`tb_plugin` submodule, which this library provides) is deprecated and scheduled for permanent removal on March 5, 2026. Users are advised to consider migrating their workflow.
  Fix: Monitor PyTorch's official channels for alternative visualization tools or migration guides. Consider using Perfetto or Chrome trace viewer for raw `trace.json` files.
• gotcha Profiling results may not be displayed in TensorBoard if the `torch.profiler` is not actively enabled and generating data in your code. Simply installing `torch-tb-profiler` is not enough to see results.
  Fix: Ensure your PyTorch code explicitly uses `torch.profiler.profile` with `tensorboard_trace_handler` to generate the necessary log files. Refer to the quickstart for an example.
• gotcha Trace files containing invalid values like 'inf' (e.g., in 'memory bandwidth (GB/s)') can cause `torch-tb-profiler` to fail when opening and visualizing the data in TensorBoard.
  Fix: This often points to a bug in the PyTorch profiler's data generation. If encountered, report an issue to the PyTorch team with a minimal reproducible example. Workaround might involve manually cleaning invalid entries if feasible.
• gotcha When profiling CUDA activities, issues like mismatched CUDA toolkit versions between your environment and PyTorch installation, or improper use of `autocast` for mixed precision, can lead to incorrect or misleading GPU profiling results and performance.
  Fix: Verify that your PyTorch installation is compatible with your system's CUDA toolkit. Ensure `torch.profiler.ProfilerActivity.CUDA` is included in your profiler activities, and that your model operations are correctly moved to the GPU. For mixed precision, ensure `torch.cuda.amp.autocast()` is used where appropriate.

Install

• pip install torch-tb-profiler Install stable version

Imports

• tensorboard_trace_handler

  from torch.profiler import tensorboard_trace_handler

  This handler is passed to `torch.profiler.profile` to generate trace files compatible with the torch-tb-profiler plugin. The plugin itself is not directly imported into user code.

Quickstart

To use `torch-tb-profiler`, first, you need to set up `torch.profiler` in your training loop to generate profiling data. Use `torch.profiler.tensorboard_trace_handler` as the `on_trace_ready` callback to save the trace files. After your script runs and generates the profiling data, launch TensorBoard from your terminal, pointing to the parent directory of your profiling logs. The plugin will automatically be available under the 'PYTORCH_PROFILER' tab in your TensorBoard UI.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity, tensorboard_trace_handler
import os

# Create a dummy model and data
model = nn.Linear(10, 10).cuda() if torch.cuda.is_available() else nn.Linear(10, 10)
optimizer = optim.SGD(model.parameters(), lr=0.01)
dummy_input = torch.randn(64, 10).cuda() if torch.cuda.is_available() else torch.randn(64, 10)

# Define log directory for TensorBoard
log_dir = "./runs/profiler_test"

# Run profiler
with profile(
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
    on_trace_ready=tensorboard_trace_handler(log_dir),
    activities=[
        ProfilerActivity.CPU,
        ProfilerActivity.CUDA if torch.cuda.is_available() else ProfilerActivity.CPU,
    ],
    record_shapes=True,
    with_stack=True
) as prof:
    for i in range(10):
        optimizer.zero_grad()
        output = model(dummy_input)
        loss = output.sum()
        loss.backward()
        optimizer.step()
        prof.step() # Advance profiler to next step

print(f"Profiling results saved to {log_dir}.\n")
print("To view in TensorBoard, run: ")
print(f"tensorboard --logdir {os.path.abspath(log_dir.split('/profiler_test')[0])}")