nnviz — Neural Network Visualiser

Live, real-time visualisation of PyTorch model training. Watch weights update, see predictions change, and understand what your network is learning — frame by frame.

Install

pip install pygame torch numpy pandas scikit-learn
# then clone or copy the nnviz/ folder into your project

(PyPI release coming soon — for now install from source.)

Quickstart

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

# 1. Any model with nn.Linear layers
model = nn.Sequential(
    nn.Linear(784, 128), nn.ReLU(),
    nn.Linear(128,  64), nn.ReLU(),
    nn.Linear( 64,  10),
)

# 2. Any DataLoader that yields (X_batch, y_batch)
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=64)

# 3. Run — a window opens, press T to start training
viz = Visualizer(
    input_shape  = (28, 28),                    # omit for tabular data
    class_names  = [str(i) for i in range(10)],
)
viz.run(model, loader, nn.CrossEntropyLoss(), optim.Adam(model.parameters()))

Controls

Visualiser parameters

What you see

┌──────────────┬─────────────┬────────────────────────────────────────┐
│  Info panel  │  Pixel grid │         Network diagram                │
│              │  (input)    │                                        │
│  ● TRAIN     │  ██░░██░░   │  ○─────○──────○──────○  ← output      │
│  Epoch  3    │  ░░████░░   │  ○  ╲  ○  ╲   ○  ╲   ○               │
│  Batch  142  │  ░░░░░░░░   │  ○  ╱  ○  ╱   ○  ╱   ○               │
│  Loss 0.312  │  ██░░░░██   │  ○─────○──────○──────○                │
│  Acc  88.3%  │             │                                        │
│              │             │  Blue lines = positive weights         │
│  Pred: 7     │             │  Red  lines = negative weights         │
│  True: 7 ✓  │             │  Brightness = magnitude                │
└──────────────┴─────────────┴────────────────────────────────────────┘

Left panel — live stats: mode, epoch, batch count, loss, accuracy, current prediction vs target.

Pixel grid — the first sample of each batch, rendered as a greyscale image (or a square grid for tabular data).

Network diagram — every inter-layer connection coloured by weight sign and magnitude. Blue = positive, red = negative, dark = near-zero (hidden so the display stays readable). The highest-confidence output neuron glows red; others stay green.

Supported datasets

Examples

python examples/mnist_example.py   # MNIST digits, image grid
python examples/iris_example.py    # Iris flowers, tabular

Roadmap

• Loss curve graph panel
• Save / load checkpoint from UI
• Activation heatmap overlay
• Multi-label classification support
• PyPI release

License

MIT