Introduction

When you train a machine learning model, numbers fly by in your terminal: loss, accuracy, learning rate, gradient norms. Staring at scrolling text is a terrible way to understand whether your model is actually learning. You need pictures — loss curves that trend downward, histograms that show whether your gradients are exploding or vanishing, and dashboards that let you compare one experiment against another.

This guide walks through the two most common tools for this job: TensorBoard, the visualization toolkit that started with TensorFlow but is now a first-class citizen in the PyTorch ecosystem, and MLflow, a full experiment-tracking platform that goes beyond plots into reproducibility and model management. Along the way, you’ll learn how to log metrics from plain PyTorch, how PyTorch Lightning makes this dramatically easier, and how to track the “boring but important” signals — gradient flow, weight distributions, learning rate schedules, images, and GPU utilization — that often reveal problems long before your loss curve does.

By the end, you should be able to:

• Set up TensorBoard and log scalars, histograms, images, and more from raw PyTorch.
• Use PyTorch Lightning’s built-in logging so you write less boilerplate.
• Diagnose training problems using gradient flow and weight histogram visualizations.
• Track GPU/compute metrics alongside your training metrics.
• Set up MLflow for experiment tracking, parameter logging, and model versioning.
• Understand when to reach for TensorBoard versus MLflow (or both).

No prior experience with either tool is assumed. Basic familiarity with PyTorch (writing a training loop, defining a model) is helpful but not strictly required — we’ll explain each snippet as we go.

What TensorBoard actually is

TensorBoard is a local web application. You run a small Python script (or a training job) that writes special log files called event files to a directory. TensorBoard reads that directory and renders the contents as interactive charts in your browser. It is not a cloud service and it does not track your experiments across machines by default — it’s a visualization layer on top of files sitting on your disk.

The key building blocks are:

• SummaryWriter: the Python object you use to write data to the log directory (often called runs/ or logs/).
• The log directory: a folder full of event files, one “run” per subfolder typically.
• The TensorBoard server: a command-line tool (tensorboard --logdir=...) that reads the log directory and serves a dashboard.

Installing TensorBoard

pip install tensorboard

If you’re using PyTorch, you don’t need TensorFlow installed — PyTorch ships its own lightweight SummaryWriter in torch.utils.tensorboard that writes TensorBoard-compatible event files.

pip install torch torchvision tensorboard

Your first scalar: logging loss

Let’s start with the simplest possible example — logging a training loss value at every step.

from torch.utils.tensorboard import SummaryWriter

# Creates (or reuses) a directory called "runs/experiment_1"
writer = SummaryWriter(log_dir="runs/experiment_1")

for step in range(100):
    fake_loss = 1.0 / (step + 1)  # pretend this is coming from training
    writer.add_scalar("Loss/train", fake_loss, step)

writer.close()

A few things to note:

• add_scalar(tag, value, step) takes a tag (a name for the metric, which can include a / to create a grouping in the UI, like Loss/train vs Loss/val), the value at this point, and the step (usually the training iteration or epoch number) — this becomes the x-axis.
• You must call writer.close() (or use it as a context manager) so all buffered data gets flushed to disk.
• Always call SummaryWriter once per run and reuse it — creating a new one every step is a common beginner mistake that fragments your logs.

To view this, run in a terminal from the same working directory:

tensorboard --logdir=runs

Then open the URL it prints (usually http://localhost:6006) in your browser. You’ll see a “Loss” section with a smooth downward curve.

Logging metrics inside a real PyTorch training loop

Here’s a more realistic example — a small classifier trained on random data, logging both training and validation loss/accuracy per epoch.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader, TensorDataset

# --- Fake data for illustration ---
X_train = torch.randn(1000, 20)
y_train = torch.randint(0, 2, (1000,))
X_val = torch.randn(200, 20)
y_val = torch.randint(0, 2, (200,))

train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32)

model = nn.Sequential(
    nn.Linear(20, 64),
    nn.ReLU(),
    nn.Linear(64, 2),
)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

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

global_step = 0
for epoch in range(10):
    # --- Training ---
    model.train()
    for xb, yb in train_loader:
        optimizer.zero_grad()
        preds = model(xb)
        loss = criterion(preds, yb)
        loss.backward()
        optimizer.step()

        # Log loss at every training step (fine-grained view)
        writer.add_scalar("Loss/train_step", loss.item(), global_step)
        global_step += 1

    # --- Validation ---
    model.eval()
    val_loss_total, correct, total = 0.0, 0, 0
    with torch.no_grad():
        for xb, yb in val_loader:
            preds = model(xb)
            val_loss_total += criterion(preds, yb).item() * xb.size(0)
            correct += (preds.argmax(dim=1) == yb).sum().item()
            total += xb.size(0)

    val_loss = val_loss_total / total
    val_acc = correct / total

    # Log epoch-level metrics (coarser view, one point per epoch)
    writer.add_scalar("Loss/val", val_loss, epoch)
    writer.add_scalar("Accuracy/val", val_acc, epoch)

    print(f"Epoch {epoch}: val_loss={val_loss:.4f}, val_acc={val_acc:.4f}")

writer.close()

Notice the pattern: fine-grained metrics (per training step, like the raw loss) are useful for spotting instability early in training, while coarse-grained metrics (per epoch, like validation accuracy) are useful for tracking overall progress. It’s common — and a good habit — to log both, using different step counters (global_step for training steps, epoch for epoch-level metrics) so the x-axes don’t get mixed up.

Comparing multiple runs

The real power of TensorBoard shows up when you compare experiments. If you point tensorboard --logdir at a parent folder containing multiple run subfolders, it overlays them automatically:

# Run 1
writer = SummaryWriter(log_dir="runs/lr_0.001")
# ... training with lr=0.001 ...

# Run 2
writer = SummaryWriter(log_dir="runs/lr_0.01")
# ... training with lr=0.01 ...

tensorboard --logdir=runs

Now the dashboard shows both lr_0.001 and lr_0.01 as separate colored lines on the same chart, letting you visually compare which learning rate converges faster. This is the foundation of hyperparameter tuning by eye.

Beyond scalars: the other TensorBoard tools

Loss curves are just the beginning. TensorBoard has dedicated views for several other data types, all attached to the same SummaryWriter.

for name, param in model.named_parameters():
    writer.add_histogram(f"weights/{name}", param, epoch)

from torchvision.utils import make_grid

images, labels = next(iter(train_loader_of_images))  # a batch of images
grid = make_grid(images[:16])  # arrange 16 images into a grid
writer.add_image("sample_batch", grid, epoch)

features = model.get_penultimate_features(X_val)  # shape: (N, D)
writer.add_embedding(features, metadata=y_val.tolist(), global_step=epoch)

dummy_input = torch.randn(1, 20)
writer.add_graph(model, dummy_input)

writer.add_hparams(
    {"lr": 1e-3, "batch_size": 32, "optimizer": "adam"},
    {"hparam/val_accuracy": val_acc, "hparam/val_loss": val_loss},
)

Logging gradient flow

Gradient flow is one of the most useful — and most overlooked — things to visualize. If gradients vanish (become near-zero) or explode (become huge) in certain layers, training will stall or diverge, often without an obvious symptom in the loss curve until it’s too late.

The simplest approach is to log the norm of each parameter’s gradient after backward():

loss.backward()

for name, param in model.named_parameters():
    if param.grad is not None:
        writer.add_scalar(f"grad_norm/{name}", param.grad.norm(), global_step)

optimizer.step()

You can also log the full gradient distribution as a histogram, which shows not just the magnitude but the shape:

for name, param in model.named_parameters():
    if param.grad is not None:
        writer.add_histogram(f"grad_hist/{name}", param.grad, global_step)

A classic and very readable way to see gradient flow across all layers at once is to compute the average absolute gradient per layer and plot it as a bar-style line — this pattern is popular enough that it’s worth writing as a reusable helper:

import numpy as np
import matplotlib.pyplot as plt

def plot_grad_flow(named_parameters):
    """Plots the average gradient magnitude per layer.
    Call this after loss.backward() and before optimizer.step()."""
    avg_grads = []
    layers = []
    for name, param in named_parameters:
        if param.requires_grad and param.grad is not None and "bias" not in name:
            layers.append(name)
            avg_grads.append(param.grad.abs().mean().item())

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(avg_grads, alpha=0.7, color="b", marker="o")
    ax.hlines(0, 0, len(avg_grads) + 1, linewidth=1, color="k")
    ax.set_xticks(range(len(layers)))
    ax.set_xticklabels(layers, rotation="vertical")
    ax.set_xlabel("Layers")
    ax.set_ylabel("Average gradient magnitude")
    ax.set_title("Gradient flow")
    fig.tight_layout()
    return fig

# Usage inside the training loop:
loss.backward()
fig = plot_grad_flow(model.named_parameters())
writer.add_figure("gradient_flow", fig, global_step)
optimizer.step()

add_figure accepts a Matplotlib figure directly, which is convenient any time you want a custom plot (not just the built-in scalar/histogram views) inside TensorBoard.

What to look for: if the average gradient magnitude is near zero for early layers but healthy for later ones, that’s a textbook vanishing-gradient signature — common in deep networks without residual connections or proper normalization. If magnitudes are enormous or NaN, you’re looking at exploding gradients, and gradient clipping (torch.nn.utils.clip_grad_norm_) is a typical fix.

Logging learning rate

Learning rate schedules are easy to log and easy to get wrong silently (e.g., a scheduler that never actually steps). Always log it:

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

for epoch in range(num_epochs):
    # ... training ...
    current_lr = optimizer.param_groups[0]["lr"]
    writer.add_scalar("LR", current_lr, epoch)
    scheduler.step()

Logging GPU and compute metrics

Understanding how fast and how efficiently your model trains is just as important as tracking accuracy — a training run that’s silently bottlenecked on data loading, or that’s about to run out of GPU memory, is a problem you want to catch early. You can log this alongside your other metrics using torch.cuda utilities (and optionally nvidia-smi/pynvml for finer-grained utilization).

import time
import torch

for epoch in range(num_epochs):
    epoch_start = time.time()

    for step, (xb, yb) in enumerate(train_loader):
        step_start = time.time()

        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        preds = model(xb)
        loss = criterion(preds, yb)
        loss.backward()
        optimizer.step()

        step_time = time.time() - step_start
        samples_per_sec = xb.size(0) / step_time

        if torch.cuda.is_available():
            mem_allocated_gb = torch.cuda.memory_allocated() / 1e9
            mem_reserved_gb = torch.cuda.memory_reserved() / 1e9
            writer.add_scalar("system/gpu_mem_allocated_gb", mem_allocated_gb, global_step)
            writer.add_scalar("system/gpu_mem_reserved_gb", mem_reserved_gb, global_step)

        writer.add_scalar("system/throughput_samples_per_sec", samples_per_sec, global_step)
        global_step += 1

    writer.add_scalar("system/epoch_time_sec", time.time() - epoch_start, epoch)

For live GPU utilization percentage (not just memory), the pynvml library (bundled with nvidia-ml-py) gives you access to the same data nvidia-smi shows:

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)  # GPU 0

util = pynvml.nvmlDeviceGetUtilizationRates(handle)
writer.add_scalar("system/gpu_util_percent", util.gpu, global_step)
writer.add_scalar("system/gpu_mem_util_percent", util.memory, global_step)

A GPU utilization consistently below ~70-80% during training often means your data loading (or CPU-side preprocessing) is the bottleneck, not your model’s compute — a useful diagnostic you’d otherwise never see.

Launching and navigating TensorBoard

tensorboard --logdir=runs --port=6006

Useful flags:

• --logdir: the parent directory containing your runs.
• --port: change the port if 6006 is taken.
• --reload_interval: how often (in seconds) TensorBoard rescans the log directory for new data (useful for watching a live run).

Inside the browser UI, the left sidebar lets you toggle which runs are visible, smooth noisy curves with a slider, and switch between tabs: Scalars, Images, Graphs, Distributions, Histograms, HParams, and Embeddings (Projector).

Why Lightning?

Everything above works, but you’ll notice a pattern: every training loop needs the same boilerplate — moving tensors to the right device, zeroing gradients, stepping the optimizer, computing epoch averages, calling writer.add_scalar in a dozen places. PyTorch Lightning is a lightweight wrapper around PyTorch that removes this boilerplate by organizing your code into a LightningModule, and — crucially for this guide — it comes with a unified logging API that automatically works with TensorBoard, MLflow, and several other backends without changing your logging calls.

pip install lightning

A minimal LightningModule with logging

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import lightning as L

class SimpleClassifier(L.LightningModule):
    def __init__(self, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()  # logs lr, etc. automatically as hyperparameters
        self.model = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 2),
        )

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        xb, yb = batch
        preds = self(xb)
        loss = F.cross_entropy(preds, yb)

        # This one line replaces writer.add_scalar + manual step tracking.
        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        xb, yb = batch
        preds = self(xb)
        loss = F.cross_entropy(preds, yb)
        acc = (preds.argmax(dim=1) == yb).float().mean()

        self.log("val_loss", loss, on_epoch=True, prog_bar=True)
        self.log("val_acc", acc, on_epoch=True, prog_bar=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)


# --- Data ---
X_train = torch.randn(1000, 20)
y_train = torch.randint(0, 2, (1000,))
X_val = torch.randn(200, 20)
y_val = torch.randint(0, 2, (200,))

train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32)

# --- Train ---
model = SimpleClassifier(lr=1e-3)
trainer = L.Trainer(max_epochs=10, log_every_n_steps=10)
trainer.fit(model, train_loader, val_loader)

That’s the entire training script. By default, Lightning creates a lightning_logs/ directory and uses TensorBoard as the logger automatically — no SummaryWriter needed. Run:

tensorboard --logdir=lightning_logs

and you’ll see the same kind of dashboard as before, generated with far less code.

Understanding self.log()

The single method self.log(name, value, ...) is the heart of Lightning’s logging system. Key arguments:

• on_step: log this value at every batch/step.
• on_epoch: aggregate (by default, averaged) this value over the epoch and log once per epoch.
• prog_bar: also show this value in the terminal progress bar, for a quick sanity check while training.
• logger: whether to send this to the configured logger(s) at all (default True).

You can log a whole dictionary at once with self.log_dict:

self.log_dict(
    {"val_loss": loss, "val_acc": acc, "val_f1": f1},
    on_epoch=True,
    prog_bar=True,
)

Explicitly choosing a logger (and using multiple at once)

By default Lightning uses TensorBoardLogger, but you can be explicit, configure it, or swap in a different backend entirely — including MLflow, which we’ll set up properly in Part 3.

from lightning.pytorch.loggers import TensorBoardLogger, MLFlowLogger

tb_logger = TensorBoardLogger(save_dir="logs", name="my_experiment")

mlf_logger = MLFlowLogger(
    experiment_name="my_experiment",
    tracking_uri="file:./mlruns",
)

trainer = L.Trainer(
    max_epochs=10,
    logger=[tb_logger, mlf_logger],  # log to both simultaneously!
)

This is one of Lightning’s most useful features for this guide’s purpose: you write self.log(...) once in your model code, and it fans out to every configured logger automatically. You don’t need separate TensorBoard and MLflow logging calls scattered through your training loop.

Logging gradient flow, weights, and other custom data in Lightning

Because self.log() is designed for scalar metrics, richer data (histograms, figures, images) is logged through hooks — methods you override that Lightning calls at specific points in training — using self.logger.experiment, which gives you direct access to the underlying TensorBoard SummaryWriter (or MLflow client, etc.).

class SimpleClassifier(L.LightningModule):
    # ... __init__, forward, training_step, configure_optimizers as before ...

    def on_before_optimizer_step(self, optimizer):
        # Called after backward(), before optimizer.step() — the ideal place
        # to inspect gradients.
        if self.trainer.global_step % 50 == 0:  # don't log every single step
            for name, param in self.named_parameters():
                if param.grad is not None:
                    grad_norm = param.grad.norm()
                    self.log(f"grad_norm/{name}", grad_norm, on_step=True)

                    # Access the raw TensorBoard writer for histograms
                    if isinstance(self.logger, L.pytorch.loggers.TensorBoardLogger):
                        self.logger.experiment.add_histogram(
                            f"grad_hist/{name}", param.grad, self.trainer.global_step
                        )

    def on_train_epoch_end(self):
        # Log weight histograms once per epoch
        if isinstance(self.logger, L.pytorch.loggers.TensorBoardLogger):
            for name, param in self.named_parameters():
                self.logger.experiment.add_histogram(
                    f"weights/{name}", param, self.current_epoch
                )

Lightning also has a built-in shortcut for a common special case — tracking whether gradients are exploding — via the Trainer’s gradient clipping arguments:

trainer = L.Trainer(
    max_epochs=10,
    gradient_clip_val=1.0,       # clip gradients to prevent explosions
    gradient_clip_algorithm="norm",
)

Lightning automatically logs the total gradient norm (as grad_2.0_norm_total) when gradient_clip_val is set, no manual code required.

Logging learning rate automatically

Rather than manually pulling optimizer.param_groups[0]["lr"], use Lightning’s built-in callback:

from lightning.pytorch.callbacks import LearningRateMonitor

lr_monitor = LearningRateMonitor(logging_interval="epoch")
trainer = L.Trainer(max_epochs=10, callbacks=[lr_monitor])

This automatically logs the learning rate of every optimizer/scheduler you configure, at every epoch (or step, if you set logging_interval="step").

Logging images in Lightning

def on_train_epoch_end(self):
    sample_images = self.get_sample_batch()  # (N, C, H, W) tensor
    grid = torchvision.utils.make_grid(sample_images)
    self.logger.experiment.add_image("samples", grid, self.current_epoch)

Logging GPU/compute metrics in Lightning

Lightning ships a built-in callback for exactly this:

from lightning.pytorch.callbacks import DeviceStatsMonitor

trainer = L.Trainer(
    max_epochs=10,
    callbacks=[DeviceStatsMonitor()],
)

DeviceStatsMonitor automatically logs GPU utilization, memory allocated/reserved, and other device stats at every step, to whichever logger(s) you’ve configured — no manual pynvml calls needed.

What MLflow is, and how it differs from TensorBoard

TensorBoard is primarily a visualization tool: point it at event files, get charts. MLflow is a broader experiment tracking and lifecycle management platform. It still gives you charts, but it also gives you:

• A structured way to log parameters (hyperparameters), metrics (time-series values, like TensorBoard scalars), and artifacts (files — model checkpoints, plots, datasets).
• A Tracking Server with a UI, backed by a database or filesystem, that stores every run’s parameters/metrics/artifacts in a queryable way — so you can filter and sort runs, not just eyeball overlaid line charts.
• A Model Registry for versioning trained models, tagging them (staging, production), and managing their lifecycle — something TensorBoard has no concept of at all.
• First-class support for team collaboration: a shared MLflow tracking server means everyone on a team sees the same experiment history, rather than everyone having their own local runs/ folder.

A simple way to think about it: TensorBoard answers “what did my loss curve look like?” MLflow answers that too, but also “which of my 40 experiments had the best validation accuracy, what hyperparameters did it use, and where’s the model file?”

Installing and running MLflow

pip install mlflow

MLflow needs a tracking URI — where to store run data. The simplest option, especially while learning, is a local folder:

mlflow ui --backend-store-uri file:./mlruns

This starts a web UI (usually at http://localhost:5000) reading from a local mlruns/ directory that MLflow creates automatically when you start logging.

For a shared/team setup, you’d instead run a proper tracking server backed by a database:

mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --default-artifact-root ./mlruns \
  --host 0.0.0.0 \
  --port 5000

For this guide, the local filesystem option is enough to learn the concepts.

Logging your first MLflow run

import mlflow

mlflow.set_tracking_uri("file:./mlruns")
mlflow.set_experiment("my_first_experiment")

with mlflow.start_run():
    mlflow.log_param("lr", 1e-3)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("optimizer", "adam")

    for epoch in range(10):
        fake_loss = 1.0 / (epoch + 1)
        fake_acc = 1 - fake_loss
        mlflow.log_metric("loss", fake_loss, step=epoch)
        mlflow.log_metric("accuracy", fake_acc, step=epoch)

Key API pieces:

• mlflow.set_experiment(name): groups runs under a named experiment (like a project folder).
• mlflow.start_run(): a context manager — everything logged inside it belongs to one run, and the run is automatically marked “finished” (or “failed”, if an exception is raised) when the block exits.
• mlflow.log_param(key, value): logs a single, unchanging hyperparameter.
• mlflow.log_metric(key, value, step): logs a time-series value — directly analogous to writer.add_scalar.

Run mlflow ui and open the browser to see this run listed, with its parameters and a metric chart.

A full PyTorch training loop with MLflow

import mlflow
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

mlflow.set_tracking_uri("file:./mlruns")
mlflow.set_experiment("classifier_experiments")

X_train = torch.randn(1000, 20)
y_train = torch.randint(0, 2, (1000,))
X_val = torch.randn(200, 20)
y_val = torch.randint(0, 2, (200,))

train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32)

lr = 1e-3
batch_size = 32
num_epochs = 10

with mlflow.start_run(run_name="baseline_run"):
    mlflow.log_params({"lr": lr, "batch_size": batch_size, "num_epochs": num_epochs})

    model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 2))
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    global_step = 0
    for epoch in range(num_epochs):
        model.train()
        for xb, yb in train_loader:
            optimizer.zero_grad()
            loss = criterion(model(xb), yb)
            loss.backward()
            optimizer.step()
            mlflow.log_metric("train_loss_step", loss.item(), step=global_step)
            global_step += 1

        model.eval()
        correct, total, val_loss_total = 0, 0, 0.0
        with torch.no_grad():
            for xb, yb in val_loader:
                preds = model(xb)
                val_loss_total += criterion(preds, yb).item() * xb.size(0)
                correct += (preds.argmax(dim=1) == yb).sum().item()
                total += xb.size(0)

        mlflow.log_metric("val_loss", val_loss_total / total, step=epoch)
        mlflow.log_metric("val_accuracy", correct / total, step=epoch)

    # Log the final trained model itself as an artifact
    mlflow.pytorch.log_model(model, artifact_path="model")

mlflow.pytorch.log_model is worth calling out: it saves the model (weights + enough metadata to reload it) directly into the run’s artifact storage, so the model file lives alongside the metrics and parameters that produced it — no more wondering which .pt file in a folder corresponds to which experiment.

Logging gradient flow, weights, and other artifacts in MLflow

MLflow’s log_metric is scalar-only, so for anything richer (histograms, gradient-flow figures, images), the pattern is: build the visualization with a library like Matplotlib, save it to a file, and log that file as an artifact.

import matplotlib.pyplot as plt
import mlflow

def log_grad_flow_figure(named_parameters, step):
    avg_grads, layers = [], []
    for name, param in named_parameters:
        if param.requires_grad and param.grad is not None and "bias" not in name:
            layers.append(name)
            avg_grads.append(param.grad.abs().mean().item())

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(avg_grads, marker="o")
    ax.set_xticks(range(len(layers)))
    ax.set_xticklabels(layers, rotation="vertical")
    ax.set_title(f"Gradient flow (step {step})")
    fig.tight_layout()

    path = f"grad_flow_step_{step}.png"
    fig.savefig(path)
    plt.close(fig)
    mlflow.log_artifact(path, artifact_path="gradient_flow")

# Inside the training loop, after loss.backward():
if global_step % 100 == 0:
    log_grad_flow_figure(model.named_parameters(), global_step)

You can also log a numeric summary (gradient norm per layer) as regular metrics, which does give you a live-updating chart in the MLflow UI:

for name, param in model.named_parameters():
    if param.grad is not None:
        mlflow.log_metric(f"grad_norm_{name.replace('.', '_')}", param.grad.norm().item(), step=global_step)

(Note: MLflow metric keys can’t contain certain characters like ., hence the replace above — a small gotcha worth knowing about early.)

Logging GPU/compute metrics in MLflow

MLflow has a built-in system metrics logger that can automatically track CPU/GPU utilization and memory during a run, no manual pynvml code required:

import mlflow

mlflow.enable_system_metrics_logging()

with mlflow.start_run():
    # ... training loop ...
    pass

This periodically logs metrics like system/gpu_0_utilization_percentage and system/gpu_0_memory_usage_percentage in the background for the duration of the run. It requires the psutil and pynvml packages (pip install psutil pynvml) to pick up CPU and GPU stats respectively.

MLflow with PyTorch Lightning

Recall from Part 2 that Lightning has a built-in MLFlowLogger. This means your self.log(...) calls “just work” with MLflow, with zero changes to your LightningModule:

from lightning.pytorch.loggers import MLFlowLogger

mlf_logger = MLFlowLogger(
    experiment_name="lightning_mlflow_demo",
    tracking_uri="file:./mlruns",
    log_model=True,  # automatically logs the final model checkpoint as an MLflow artifact
)

trainer = L.Trainer(
    max_epochs=10,
    logger=mlf_logger,
    callbacks=[DeviceStatsMonitor(), LearningRateMonitor(logging_interval="epoch")],
)
trainer.fit(model, train_loader, val_loader)

Everything from Part 2 — gradient norms, learning rate, GPU stats — flows into MLflow automatically, exactly as it did for TensorBoard, because it’s all going through the same self.log() interface.

The MLflow UI: comparing runs

Start the UI:

mlflow ui --backend-store-uri file:./mlruns

In the browser, you’ll see a table of all runs in an experiment, one row per run, with columns for every logged parameter and metric. This is the biggest practical difference from TensorBoard’s UI: you can sort and filter runs by their final metric values or hyperparameters (e.g., “show me every run with lr < 0.001 sorted by val_accuracy descending”) rather than visually comparing overlaid line charts. You can also select multiple runs and get an automatic parallel-coordinates plot or scatter plot comparing hyperparameters against outcomes — extremely useful once you’ve run more than a handful of experiments.

The Model Registry

Once you’ve found a run you’re happy with, MLflow lets you promote its logged model into a registry — a separate, versioned catalog of models with lifecycle stages.

result = mlflow.register_model(
    model_uri="runs:/<RUN_ID>/model",
    name="my_classifier",
)

Or directly during logging:

mlflow.pytorch.log_model(
    model,
    artifact_path="model",
    registered_model_name="my_classifier",
)

From the UI (or API), you can then tag versions as Staging or Production, and load the “current production model” by name rather than by a specific run ID or file path:

import mlflow.pytorch

model = mlflow.pytorch.load_model("models:/my_classifier/Production")

This is the piece with no TensorBoard equivalent at all — TensorBoard has no concept of “the model,” only the metrics you chose to visualize.

Minimal cheat sheet — raw PyTorch + TensorBoard

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("runs/exp1")
writer.add_scalar("loss", loss_value, step)
writer.add_histogram("weights/layer1", param, step)
writer.add_image("samples", img_grid, step)
writer.close()

Minimal cheat sheet — raw PyTorch + MLflow

import mlflow
mlflow.set_experiment("exp1")
with mlflow.start_run():
    mlflow.log_param("lr", 1e-3)
    mlflow.log_metric("loss", loss_value, step=step)
    mlflow.log_artifact("figure.png")
    mlflow.pytorch.log_model(model, "model")

Minimal cheat sheet — PyTorch Lightning

self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
# Trainer picks the backend:
trainer = L.Trainer(logger=[TensorBoardLogger("logs"), MLFlowLogger(tracking_uri="file:./mlruns")])