" MicromOne: How to Use TensorBoard with PyTorch

Pagine

How to Use TensorBoard with PyTorch

TensorBoard allows you to visualize:

  • Training and validation loss

  • Accuracy and other custom metrics

  • Images, model graphs, and embeddings

  • Comparisons between multiple training runs

Instead of relying only on printed logs, TensorBoard provides an interactive dashboard that makes it much easier to track what is happening during training.

Installing TensorBoard

If you are using PyTorch, TensorBoard is usually installed automatically. If not, you can install it manually:

pip install tensorboard

Make sure you also have PyTorch installed:

pip install torch torchvision

Launching TensorBoard

TensorBoard is started from the command line. First, choose a directory where your training logs will be stored (for example, runs/). Then run:

tensorboard --logdir=runs

After launching, TensorBoard will print a local URL, typically:

http://localhost:6006

Open this URL in your browser to access the TensorBoard dashboard.

Using SummaryWriter in PyTorch

PyTorch provides TensorBoard integration through the SummaryWriter class, which is located in:

from torch.utils.tensorboard import SummaryWriter

The SummaryWriter is responsible for writing event files that TensorBoard reads and visualizes.

Creating a SummaryWriter

writer = SummaryWriter(log_dir="runs/experiment_1")

Each experiment should ideally have its own log directory. This makes it easy to compare multiple runs in TensorBoard.

Logging Scalars (Loss, Accuracy, Metrics)

The most common use of TensorBoard is logging scalar values such as loss and accuracy during training.

Example: Logging Training Loss

for epoch in range(num_epochs):
    train_loss = 0.0
    
    # training loop
    for inputs, labels in dataloader:
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    avg_loss = train_loss / len(dataloader)
    writer.add_scalar("Loss/train", avg_loss, epoch)

Logging Accuracy

writer.add_scalar("Accuracy/train", accuracy, epoch)

In TensorBoard, these values will appear as smooth, interactive plots that update as training progresses.

Logging Images

TensorBoard can also display images, which is especially useful for computer vision tasks.

Example: Logging a Batch of Images

from torchvision.utils import make_grid

images, labels = next(iter(dataloader))
img_grid = make_grid(images)

writer.add_image("Training Images", img_grid)

You can view these images in the Images tab of TensorBoard.

Logging Figures

Sometimes you may want to log custom plots, such as confusion matrices or matplotlib charts. TensorBoard supports this using the add_figure method.

Example: Logging a Matplotlib Figure

import matplotlib.pyplot as plt

fig, ax = plt.subplots()
ax.plot([1, 2, 3], [4, 5, 6])
ax.set_title("Sample Plot")

writer.add_figure("My Figure", fig, global_step=epoch)

This is useful for visualizing advanced statistics or custom evaluation results.

Closing the Writer

After training is complete, always close the SummaryWriter to ensure all data is written to disk:

writer.close()

Best Practices

  • Use clear tag names (e.g., Loss/train, Loss/val)

  • Create separate log directories for different experiments

  • Log metrics at consistent intervals

  • Avoid logging too frequently to reduce disk usage