dead-relu-viz

Visualize dead ReLU neurons directly on a PyTorch model's architecture graph.

A dead neuron is one whose activation never leaves zero across an entire dataset — a common failure mode after a bad weight init or too high a learning rate. dead-relu-viz propagates your data through your model, measures which Conv2d/Linear units never activate, and overlays the result on the model's architecture graph (built with torchview), so you can spot where dead neurons are concentrated at a glance.

Works with any PyTorch nn.Module and any DataLoader — not tied to a specific model or dataset.

Installation

pip install dead-relu-viz

Graphviz must also be installed on your system (it's what torchview uses to render the graph):

apt install graphviz      # Debian/Ubuntu
brew install graphviz     # macOS

Usage

In a Jupyter notebook:

from dead_relu_viz import visualization

# model: any PyTorch nn.Module
# data_loader: any DataLoader
visualization(model, data_loader, with_activation=True, dpi=300)

This displays the model's architecture graph inline, with each Conv2d/Linear layer that contains dead neurons highlighted in red and annotated with the count and percentage of dead units.

Call it with with_activation=False (the default) to just render the plain architecture graph, without running any data through the model.

Arguments

• model — any PyTorch model (nn.Module).
• data — a DataLoader used to measure activations.
• with_activation — False (default): plain torchview graph. True: graph annotated with dead neurons in red.
• device — device used to run the model when measuring activations (defaults to CUDA if available, else CPU).
• dpi — if set, overrides the graph's rendering resolution (e.g. dpi=300).
• exclude_from_dead_check — layer names (as in model.named_modules()) to skip when coloring dead neurons — still measured and printed, just not flagged red. Useful for layers where "dead" isn't meaningful (e.g. a layer followed by a sigmoid instead of a ReLU).

Example

See cnn-activations for two worked examples (a small CNN trained on CIFAR-10, and GMIC, a breast-cancer screening model) using this tool.

License

MIT