Kahibaro
Discord Login Register

6.6 Logging with TensorBoard

What TensorBoard Gives You During Training

TensorBoard is a lightweight dashboard for viewing training progress while your model runs. In a beginner friendly PyTorch workflow, it is mainly used to track scalars like training loss, validation loss, accuracy, and learning rate over time. It can also display parameter histograms, gradient histograms, images, and embeddings, but the most common and most useful starting point is scalar logging, because it immediately answers whether training is improving and whether it is overfitting.

The typical pattern is simple: you create a log writer, you write values to it at regular steps, you close it at the end, then you launch TensorBoard pointing at the log directory. The only thing you must decide is what you consider a step. Most projects use either a global step that increments every optimization update, or an epoch based step that increments once per epoch. Global step is usually more informative because it aligns directly with gradient updates.

Installing and Importing the Writer

In PyTorch, TensorBoard logging is provided through torch.utils.tensorboard. If you installed PyTorch normally, you often already have what you need. If you see an import error, installing the tensorboard package usually resolves it.

In code, the entry point is SummaryWriter, which creates an event file inside a directory you choose.

python
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir="runs/exp1")

The log_dir should be unique per experiment run, otherwise multiple runs will mix together and your plots become confusing. A common approach is to include a timestamp or a short run name that encodes key hyperparameters.

Do not reuse the same log_dir across different experiments unless you explicitly want TensorBoard to merge them. Use a fresh directory per run to keep comparisons clean.

Logging Scalars: Loss, Metrics, and Learning Rate

The core method is writer.add_scalar(tag, value, global_step). The tag is a string that defines the series name, the value is a Python number, and global_step is the x axis.

A common convention is to namespace tags so related plots group together, for example Loss/train, Loss/val, Accuracy/train, Accuracy/val.

python
global_step = 0
for epoch in range(num_epochs):
    model.train()
    for xb, yb in train_loader:
        optimizer.zero_grad()
        preds = model(xb)
        loss = criterion(preds, yb)
        loss.backward()
        optimizer.step()
        writer.add_scalar("Loss/train", loss.item(), global_step)
        lr = optimizer.param_groups[0]["lr"]
        writer.add_scalar("LR", lr, global_step)
        global_step += 1

Validation is typically logged once per epoch, because running validation every batch can be expensive and noisy. In that case, you still log against a step value, either the current global_step or the epoch number. Using global_step makes training and validation curves line up on the same horizontal scale.

python
model.eval()
val_loss = 0.0
correct = 0
count = 0
with torch.no_grad():
    for xb, yb in val_loader:
        preds = model(xb)
        loss = criterion(preds, yb)
        val_loss += loss.item() * xb.size(0)
        # example for classification
        predicted = preds.argmax(dim=1)
        correct += (predicted == yb).sum().item()
        count += xb.size(0)
val_loss /= count
val_acc = correct / count
writer.add_scalar("Loss/val", val_loss, global_step)
writer.add_scalar("Accuracy/val", val_acc, global_step)

Always log Python numbers, not tensors that still require gradients. Use loss.item() or float(value) so logging does not accidentally keep parts of the computation graph alive and waste memory.

Choosing a Step Convention That Stays Consistent

TensorBoard plots are only as readable as your step definition. If you log training per batch and validation per epoch, you should still use a single consistent step axis. A practical rule is to use global optimization step for everything, because it is monotonic and comparable across runs even if you change batch size or number of batches.

If you prefer epoch based logging, then log both training and validation using epoch as the step, which means you should aggregate training metrics across the whole epoch before writing them.

Do not mix different meanings of global_step under the same tag. If Loss/train uses batch steps, then keep it that way for the whole run.

Logging Histograms: Weights and Gradients

Once scalar curves look reasonable, histograms can help you diagnose issues like exploding weights or dead activations. You can log parameter values with add_histogram. Doing this every step is expensive, so doing it every few epochs is common.

python
if epoch % 5 == 0:
    for name, param in model.named_parameters():
        writer.add_histogram(f"Weights/{name}", param.detach().cpu().numpy(), global_step)
        if param.grad is not None:
            writer.add_histogram(f"Grads/{name}", param.grad.detach().cpu().numpy(), global_step)

Histograms are most useful when you know what you are looking for. For example, gradients that become extremely concentrated near zero across layers can hint at vanishing gradients, and gradients with very large magnitude spikes can hint at instability.

Logging Images: Predictions and Inputs (Optional but Powerful)

For vision tasks, logging a small batch of input images and model outputs can catch data pipeline mistakes quickly, such as wrong normalization or swapped channels. TensorBoard accepts images in CHW format per image and values typically in $[0, 1]$ or $[0, 255]$ depending on how you log.

A simple pattern is to log a fixed batch from the validation set every few epochs so you can visually compare changes over time.

python
import torchvision
# xb: (N, C, H, W)
grid = torchvision.utils.make_grid(xb[:16], nrow=4, normalize=True)
writer.add_image("Images/val_batch", grid, global_step)

If you also want to log predicted labels, you usually include them in the tag name or write a small text summary separately. Keep it minimal at first, since images alone already catch many issues.

Logging Text: Short Summaries and Config

TensorBoard can also store text. This is useful for recording hyperparameters, dataset version, or a short run description.

python
writer.add_text("Run/notes", "Baseline MLP, lr=1e-3, batch=128", 0)

If you want structured hyperparameter logging, PyTorch also supports add_hparams, which creates a compact table view of hyperparameters and final metrics. This is best used when you have a consistent set of hyperparameters across many runs.

Where to Put Logging in a Clean Training Loop

A clean separation is to compute metrics in your training and validation code, then log them immediately after you have the numbers. Training loss is often logged per step, while epoch level summaries like mean training loss, mean validation loss, and validation accuracy are logged once per epoch.

At the end of training, always close the writer so all events flush to disk.

python
writer.close()

If you stop a run abruptly, some recent events may not be written to disk yet. Closing the writer at the end, and occasionally flushing during very long runs, reduces the chance of losing logs.

Launching TensorBoard and Viewing Runs

TensorBoard reads event files from a directory. If your logs are in runs/, you run TensorBoard pointing there and open the provided local URL in your browser.

bash
tensorboard --logdir runs

Each subdirectory under runs/ appears as a separate run. Naming those directories well makes comparison easy, especially when you overlay training curves from multiple runs.

Minimal Checklist for Reliable TensorBoard Logging

Use a unique log directory per experiment, use a consistent step definition, log only plain numbers for scalars, and remember to close the writer so data is flushed to disk.

Views: 77

Comments

Please login to add a comment.

Don't have an account? Register now!