pysalty 0.1.0

PySalty(Python Structural Architecture Layout Tool), publication-ready PyTorch model visualization library

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

---

Installation

Install from PyPI:

pip install pysalty

---

Quick Start

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

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

model = SimpleNet()

Visualize(
    model,
    theme=PaperTheme,
    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.

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.

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=11,
    TitleFontSize=13,
    FontWeight="bold",

    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.