Prerequisites: Graphviz

The visualization libraries hiddenlayer and torchviz both rely on Graphviz for generating graph images. Download the installer from graphviz.org/download (any recent stable version, roughly 5 MB). During installation, make sure to add Graphviz to the system PATH.

After Graphviz is installed, install the Python packages:

pip install hiddenlayer torchviz

Visualizing Network Architecture

Using hiddenlayer

Assume we define a simple CNN for MNIST:

import torch
import torch.nn as nn

class SimpleMNIST(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(16 * 4 * 4, 64),
            nn.ReLU(),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

net = SimpleMNIST()

To render the structure, provide a dummy input tensor matching the expected shape:

import hiddenlayer as hl

vis_graph = hl.build_graph(net, torch.zeros(1, 1, 28, 28))
vis_graph.theme = hl.graph.THEMES["blue"].copy()
vis_graph.save("simple_mnist_hiddenlayer.png")

The output is a vertically arranged block diagram. You can switch themes or customise node colours. For deeper architectures like ResNet, the graph becomes long but remains readable.

Using torchviz

Torchviz hooks into autograd to construct a backward graph, so the input must require gradients:

from torchviz import make_dot

x = torch.randn(1, 1, 28, 28, requires_grad=True)
out = net(x)
dot = make_dot(out, params=dict(net.named_parameters()))
dot.format = "png"
dot.render("simple_mnist_torchviz")

By default this produces a graph in the current directory. The layout is often more compact than hiddenlayer, but image resolution may be low; you can tweak dot.format and dot.engine if needed. For large models, rendering can be computationally heavy.

Both approaches are non‑intrusive: just instantiate the model, feed a dummy tensor, and call the visualization function – no changes to your original model class.

Monitoring Training with TensorBoard

TensorBoard integration via tensorboardX allows you to log scalars, images, histograms, and even the model graph during training. First install the package:

pip install tensorboardX

Basic setup

from tensorboardX import SummaryWriter

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

The following snippet demonstrates logging during a training loop, using a ResNet‑like model on CIFAR‑10:

import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Data preparation
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

train_set = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform_train)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)

# Model, loss, optimizer
model = ResNet(num_classes=10)  # assume ResNet is defined elsewhere
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

# Training loop
for epoch in range(5):
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        global_step = epoch * len(train_loader) + i + 1

        if i % 100 == 0:
            print(f"Epoch [{epoch+1}/5], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}")

        # Logging
        writer.add_scalar("Loss/train", loss.item(), global_step)
        # Log a sample image (every 100 steps)
        if i % 100 == 0:
            writer.add_image("Images/train_sample", images[0], global_step)
        # Log parameter histograms (every 200 steps)
        if i % 200 == 0:
            for name, param in model.named_parameters():
                writer.add_histogram(name, param.data.cpu().numpy(), global_step)

# Log model graph once (using a dummy batch)
dummy_input = torch.randn(1, 3, 32, 32)
writer.add_graph(model, dummy_input)

writer.close()

After training, launch TensorBoard from the terminal:

tensorboard --logdir=runs

Open the printed URL in a browser to explore loss curves, image previews, distribution of weights, and the computation graph.

Recording evaluation metrics

During validation you can log accuracy or a confusion matrix:

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        correct += (predicted == labels).sum().item()
        total += labels.size(0)
writer.add_scalar("Accuracy/test", correct / total, epoch)

Use add_scalar for any numeric metric and add_image for visualising model outputs or misclassified examples.

Cloud‑native Logging with Weights & Biases

Weights & Biases (wandb) offers similar functionality with additional experiment management features. Install it and login:

pip install wandb
wandb login   # requires an API key from wandb.ai

Initialize a project before training:

import wandb

wandb.init(project="resnet-cifar10", config={
    "learning_rate": 0.01,
    "epochs": 5,
    "batch_size": 128,
})
config = wandb.config

Within the training loop you can log values using wandb.log():

for epoch in range(config.epochs):
    for i, (images, labels) in enumerate(train_loader):
        # … forward / backward / optimize …
        wandb.log({
            "train/loss": loss.item(),
            "train/step": global_step,
        }, step=global_step)
        if i % 100 == 0:
            wandb.log({"train/sample_image": wandb.Image(images[0])})

# Optionally watch model gradients and parameters
wandb.watch(model, log="all")

# Final logging
wandb.log({"test/accuracy": correct / total})
wandb.finish()

Wandb automatically captures system metrics (CPU/GPU utilisation, memory) and provides a dashobard for comparing runs, viewing live logs, and sharing results. The initialisation and logging steps are straightforward and do not interfere with training.

Choosing a Tool

• Network architecture diagrams: hiddenlayer produces clean static images, while torchviz renders autograd graphs that show operations. Use either for quick visual checks of your model structure.
• Training progress: TensorBoard is a mature, locally run tool that integrates tight with PyTorch via tensorboardX. Weights & Biases adds cloud storage, run comparison, and system monitoring. Both require only a few extra lines in your training script.

Any of these tools can be added to a existing PyTorch project with minimal refactoring. Experiment to find the combination that best fits your workflow.