PySalty (Python Structural Architecture Layout Tool)

PySalty is a lightweight Python library for visualizing PyTorch neural network architectures as clean, publication ready diagrams.

It is designed for:

• research papers
• presentations
• documentation
• understanding complex models

PySalty focuses on clarity and aesthetics rather than training, execution, or graph theory.

---

Features

• Automatic PyTorch model parsing
• Hierarchical visualization of nn.Sequential blocks
• Intelligent edge routing that avoids overlapping modules
• Multiple built in visual themes
• Fully customizable colors, fonts, and layout
• High resolution PNG export
• Zero dependency on model execution
• No arrows or ports for clean publication ready figures

---

Installation

Install from PyPI:

pip install pysalty

---

Quick Start

import torch.nn as nn
from salt import Visualize, PaperTheme, DarkTheme

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.Stem = nn.Sequential(
            nn.Conv2d(3, 64, 3),
            nn.ReLU(),
            nn.Conv2d(3, 64, 3),
            nn.ReLU(),
            nn.Conv2d(3, 64, 3),
        )
        self.Head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(64, 10),
            nn.ReLU(),
            nn.Conv2d(3, 64, 3),
            nn.ReLU(),
            nn.Conv2d(3, 64, 3),
        )
        self.Custom_Layer = nn.Sequential(nn.Flatten(), nn.AdaptiveAvgPool1d(1), nn.Flatten())

model = SimpleNet()

Visualize(
    model,
    theme=DarkTheme,
    save_path="model.png"
)

This produces a clean architecture diagram suitable for papers or slides.

---

How model parsing works

PySalty parses only the top level children of your PyTorch model.

Example structure:

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.Stem = ...
        self.Stage1 = ...
        self.Stage2 = ...
        self.Head = ...

Each top level attribute becomes a block in the diagram.

---

Containers nn.Sequential

If a top level module is an nn.Sequential, PySalty will:

• Draw it as a container box
• Draw each sub layer inside it
• Automatically connect internal layers vertically

Example:

self.Stem = nn.Sequential(
    ConvBNAct(...),
    ConvBNAct(...),
)

This becomes a container labeled Stem with child blocks inside.

---

Non sequential modules

If a module is not an nn.Sequential, it is drawn as a single block labeled by its class name.

---

Connections

• Blocks are connected in the order they appear in the model definition
• Routing automatically avoids overlapping other blocks
• Long connections are lifted above the diagram when needed
• No manual configuration is required

---

Built in Themes

PySalty includes several ready to use themes defined in salt.theme.

DarkTheme

Dark background with soft purple highlights. Suitable for presentations and screen viewing.

from salt import DarkTheme

LightTheme

Clean white background with subtle gray outlines. Suitable for documentation.

from salt import LightTheme

PaperTheme

Pure black and white styling with sharp corners and thinner edges. Recommended for research papers and LaTeX figures.

from salt import PaperTheme

DraculaTheme

Dark purple background with neon accents inspired by the Dracula color palette.

from salt import DraculaTheme

---

Creating your own theme

You can fully customize the appearance by creating a Theme object. All styling options are defined in salt.theme.Theme.

from salt import Theme

MyTheme = Theme(
    Background="#0B1020",
    Primary="#00E5FF",
    Accent="#FFD54F",
    Text="#E0F7FA",
    Line="#90A4AE",
    BoxFill="#102027",
    ChildFill="#1C313A",

    FontFamily="Times New Roman",
    FontSize=12,
    TitleFontSize=13,
    FontWeight="regular",

    CornerRadius=0.3,
    EdgeWidth=2.5,
)

Then use it:

Visualize(model, theme=MyTheme)

---

Fonts

PySalty uses Matplotlib fonts.

Any font installed on your system can be used by setting:

FontFamily="Times New Roman"

Common choices for papers include Times New Roman, Computer Modern, and DejaVu Serif.

---

Output

• By default, the figure is returned and can be shown inline in Jupyter
• Use save_path to export a high resolution PNG

Visualize(model, save_path="architecture.png")

---

License

MIT License

Free for academic and commercial use.

---

Notes and Limitations

• PySalty is a visualization tool, not a graph execution engine
• Skip connections and general DAGs are not yet supported
• Only top level modules are visualized by design

These decisions are intentional to keep diagrams clean and readable.

---

Contributing

Contributions are welcome, including new themes, layout improvements, export formats, and documentation improvements.