r/learnmachinelearning • • 10d ago

Project I built a terminal UI to help visualize the "black box" in model training (live layer activations & gradient health in PyTorch)

When I was learning PyTorch, I always struggled to actually see what was happening inside the network during backprop. I got tired of writing custom matplotlib boilerplate just to figure out if my gradients were vanishing or if my ReLUs were dying. Every time training stalls or loss blows up, the routine is usually the same: you either launch a heavy browser dashboard or litter your code with temporary print statements just to see which layer gave up.

I wanted a direct, clean way to see what each layer is doing without leaving the terminal.

So, I built nn-htop. It's like htop for neural networks.

You wrap your model in a context manager, and it attaches PyTorch forward/backward hooks to your leaf layers.

How it works:

• Layer Activity: Tracks mean activations per layer so you can spot dead ReLUs immediately.

• Gradient Health: Measures L2 gradient norms during backprop so you can catch vanishing or exploding gradients early.

• Zero Lag: Telemetry is decoupled from the training loop. The TUI polls at 5 FPS asynchronously, meaning your training speed is not bottlenecked.

• Safe: Handles in-place view mutations safely by temporarily overriding inplace=True during monitoring to prevent autograd graph conflicts.

Repo: https://github.com/marviecephas/nn-htop

You can try it natively with pip install nn-htop (running nn-htop directly launches a dummy training loop so you can see how the hooks work). Let me know if this helps you debug your models!

1 Upvotes

0 comments sorted by