Skip to content

README.md

Latest commit

 

History

History
85 lines (57 loc) · 1.92 KB

README.md

File metadata and controls

85 lines (57 loc) · 1.92 KB
Raw

NetPlotPy

Visualize PyTorch neural networks as node graphs exported to PNG.

Examples

examples/001.py - custom nn.Module examples/002.py - nn.Sequential examples/004.py - layer labels for bigger layers
001 example 002 example 004 example

Install

PyPI release coming soon! For now, install locally:

git clone https://github.com/kingazm/NetPlotPy.git
cd NetPlotPy
pip install -e .

For development (includes pytest, pdoc, pre-commit):

pip install -e ".[dev]"

Quick start

import torch.nn as nn
from NetPlotPy import draw

model = nn.Sequential(
    nn.Linear(3, 3),
    nn.Linear(3, 1),
)

draw(model, "output.png")

Supported layers

Currently supported:

Layer Notes
nn.Linear Input and output nodes are drawn per feature

nn.Sequential is supported as the top-level container when all its children are supported layers.

Broader layer support (activations, conv, recurrent) is planned. Until then, unsupported layer types raise a ValueError at runtime.

Extending with custom layers

Register a parser for any unsupported layer type:

import torch.nn as nn
from NetPlotPy import register_layer_parser
from NetPlotPy.models import LayeredGraph

def parse_conv2d(module: nn.Conv2d, graph: LayeredGraph) -> int:
    graph.add_layer(module.in_channels)
    return module.out_channels

register_layer_parser(nn.Conv2d, parse_conv2d)

The parser receives the module and the graph being built. It should call graph.add_layer(n) for the input size and return the output size as an int.

Docs

Generated docs are in docs/. To regenerate:

pdoc NetPlotPy -o docs/
xdg-open docs/NetPlotPy.html   # Linux
open docs/NetPlotPy.html        # macOS

Running tests

python -m pytest tests/