visualtorch 0.2.1__tar.gz → 0.2.2__tar.gz

{visualtorch-0.2.1/visualtorch.egg-info → visualtorch-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: visualtorch
-Version: 0.2.1
+Version: 0.2.2
 Summary: Architecture visualization of Torch models
 Home-page: https://github.com/willyfh/visualtorch
 Author: Willy Fitra Hendria
@@ -23,13 +23,13 @@ License-File: LICENSE
 
 </div>
 
-**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style and graph-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
+**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
 
 **Note:** VisualTorch may not yet support complex models, but contributions are welcome!
 
 <div align="center">
 
-![VisualTorch Examples](https://github.com/willyfh/visualtorch/assets/5786636/7e2c35ea-d34d-4b92-b414-285bccb8576a)
+![VisualTorch Examples](https://github.com/willyfh/visualtorch/assets/5786636/398c3356-4de0-446b-a30b-d8ebe532d2c2)
 
 </div>
 

{visualtorch-0.2.1 → visualtorch-0.2.2}/README.md

@@ -5,13 +5,13 @@
 
 </div>
 
-**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style and graph-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
+**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating layered-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
 
 **Note:** VisualTorch may not yet support complex models, but contributions are welcome!
 
 <div align="center">
 
-![VisualTorch Examples](https://github.com/willyfh/visualtorch/assets/5786636/7e2c35ea-d34d-4b92-b414-285bccb8576a)
+![VisualTorch Examples](https://github.com/willyfh/visualtorch/assets/5786636/398c3356-4de0-446b-a30b-d8ebe532d2c2)
 
 </div>
 

visualtorch-0.2.2/pyproject.toml

@@ -0,0 +1,167 @@
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+# SETUP CONFIGURATION.                                                        #
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+# RUFF CONFIGURATION                                                          #
+[tool.ruff]
+# Enable rules
+select = [
+    "F",    # Pyflakes (`F`)
+    "E",    # pycodestyle error (`E`)
+    "W",    # pycodestyle warning (`W`)
+    "C90",  # mccabe (`C90`)
+    "I",    # isort (`I`)
+    "N",    # pep8-naming (`N`)
+    "D",    # pydocstyle (`D`)
+    "UP",   # pyupgrade (`UP`)
+    "YTT",  # flake8-2020 (`YTT`)
+    "ANN",  # flake8-annotations (`ANN`)
+    "S",    # flake8-bandit (`S`)
+    "BLE",  # flake8-blind-except (`BLE`)
+    "FBT",  # flake8-boolean-trap (`FBT`)
+    "B",    # flake8-bugbear (`B`)
+    "A",    # flake8-builtins (`A`)
+    "COM",  # flake8-commas (`COM`)
+    "CPY",  # flake8-copyright (`CPY`)
+    "C4",   # flake8-comprehensions (`C4`)
+    "DTZ",  # flake8-datatimez (`DTZ`)
+    "T10",  # flake8-debugger (`T10`)
+    "EM",   # flake8-errmsg (`EM`)
+    "FA",   # flake8-future-annotations (`FA`)
+    "ISC",  # flake8-implicit-str-concat (`ISC`)
+    "ICN",  # flake8-import-conventions (`ICN`)
+    "PIE",  # flake8-pie (`PIE`)
+    "PT",   # flake8-pytest-style (`PT`)
+    "RSE",  # flake8-raise (`RSE`)
+    "RET",  # flake8-return (`RET`)
+    "SLF",  # flake8-self (`SLF`)
+    "SIM",  # flake8-simplify (`SIM`)
+    "TID",  # flake8-tidy-imports (`TID`)
+    "TCH",  # flake8-type-checking (`TCH`)
+    "INT",  # flake8-gettext (`INT`)
+    "ARG",  # flake8-unsused-arguments (`ARG`)
+    "PTH",  # flake8-use-pathlib (`PTH`)
+    "TD",   # flake8-todos (`TD`)
+    "FIX",  # flake8-fixme (`FIX`)
+    "ERA",  # eradicate (`ERA`)
+    "PD",   # pandas-vet (`PD`)
+    "PGH",  # pygrep-hooks (`PGH`)
+    "PL",   # pylint (`PL`)
+    "TRY",  # tryceratos (`TRY`)
+    "FLY",  # flynt (`FLY`)
+    "NPY",  # NumPy-specific rules (`NPY`)
+    "PERF", # Perflint (`PERF`)
+    "RUF",  # Ruff-specific rules (`RUF`)
+    # "FURB", # refurb (`FURB`) - ERROR: Unknown rule selector: `FURB`
+    # "LOG",  # flake8-logging (`LOG`) - ERROR: Unknown rule selector: `LOG`
+]
+
+ignore = [
+    # pydocstyle
+    "D107", # Missing docstring in __init__
+    "D415", # First line should end with a period, question mark, or exclamation point
+
+    # pylint
+    "PLR0913",  # Too many arguments to function call
+    "PLR2004",  # consider replacing with a constant variable
+    "PLR0912",  # Too many branches
+    "PLR0915",  # Too many statements
+
+    # flake8-annotations
+    "ANN101",   # Missing-type-self
+    "ANN002",   # Missing type annotation for *args
+    "ANN003",   # Missing type annotation for **kwargs
+
+    # flake8-bandit (`S`)
+    "S101", # Use of assert detected.
+
+    # flake8-boolean-trap (`FBT`)
+    "FBT001",   # Boolean positional arg in function definition
+    "FBT002",   # Boolean default value in function definition
+
+    # flake8-datatimez (`DTZ`)
+    "DTZ005",    # The use of `datetime.datetime.now()` without `tz` argument is not allowed
+
+    # flake8-fixme (`FIX`)
+    "FIX002",    # Line contains TODO, consider resolving the issue
+]
+
+# Allow autofix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Exclude a variety of commonly ignored directories.
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".hg",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "venv",
+]
+
+# Same as Black.
+line-length = 120
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+# Assume Python 3.10.
+target-version = "py310"
+
+# Allow imports relative to the "src" and "tests" directories.
+src = ["visualtorch", "tests"]
+
+[tool.ruff.mccabe]
+# Unlike Flake8, default to a complexity level of 10.
+max-complexity = 15
+
+[tool.ruff.per-file-ignores]
+"tests/nightly/tools/benchmarking/test_benchmarking.py" = ["E402"]
+
+[tool.ruff.pydocstyle]
+convention = "google"
+
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+# MYPY CONFIGURATION.                                                         #
+[tool.mypy]
+ignore_missing_imports = true
+show_error_codes = true
+
+
+[[tool.mypy.overrides]]
+module = ["torch.*"]
+follow_imports = "skip"
+follow_imports_for_stubs = true
+
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+# PYTEST CONFIGURATION                                                        #
+[tool.pytest.ini_options]
+addopts = [
+    "--strict-markers",
+    "--strict-config",
+    "--showlocals",
+    "-ra",
+]
+testpaths = "tests"
+pythonpath = "visualtorch"

{visualtorch-0.2.1 → visualtorch-0.2.2}/setup.py

@@ -3,14 +3,17 @@
 # Copyright (C) 2024 Willy Fitra Hendria
 # SPDX-License-Identifier: MIT
 
+from pathlib import Path
+
 import setuptools
 
-with open("README.md", "r") as fh:
+file_path = Path("README.md")
+with file_path.open("r") as fh:
     long_description = fh.read()
 
 setuptools.setup(
     name="visualtorch",
-    version="0.2.1",
+    version="0.2.2",
     author="Willy Fitra Hendria",
     author_email="willyfitrahendria@gmail.com",
     description="Architecture visualization of Torch models",

{visualtorch-0.2.1 → visualtorch-0.2.2}/visualtorch/__init__.py

@@ -3,7 +3,8 @@
 # Copyright (C) 2024 Willy Fitra Hendria
 # SPDX-License-Identifier: MIT
 
-from visualtorch.layered import layered_view
 from visualtorch.graph import graph_view
+from visualtorch.layered import layered_view
+from visualtorch.lenet_style import lenet_view
 
-__all__ = ["layered_view", "graph_view"]
+__all__ = ["layered_view", "graph_view", "lenet_view"]

{visualtorch-0.2.1 → visualtorch-0.2.2}/visualtorch/graph.py

@@ -3,35 +3,36 @@
 # Copyright (C) 2024 Willy Fitra Hendria
 # SPDX-License-Identifier: MIT
 
-import aggdraw
-from PIL import Image
+from collections import defaultdict
 from math import ceil
-from .layer_utils import model_to_adj_matrix, add_input_dummy_layer, TARGET_OPS
-from .utils import Circle, Ellipses, get_keys_by_value, Box
+from typing import Any
+
+import aggdraw
 import numpy as np
-from typing import Optional, Dict, Any, Tuple, List
 import torch
+from PIL import Image
+
+from .utils.layer_utils import TARGET_OPS, add_input_dummy_layer, model_to_adj_matrix
+from .utils.utils import Box, Circle, Ellipses, get_keys_by_value
 
 
 def graph_view(
     model: torch.nn.Module,
-    input_shape: Tuple[int, ...],
-    to_file: Optional[str] = None,
-    color_map: Optional[Dict[Any, Any]] = None,
+    input_shape: tuple[int, ...],
+    to_file: str | None = None,
+    color_map: dict[Any, Any] | None = None,
     node_size: int = 50,
-    background_fill: Any = "white",
+    background_fill: str | tuple[int, ...] = "white",
     padding: int = 10,
     layer_spacing: int = 250,
     node_spacing: int = 10,
-    connector_fill: Any = "gray",
+    connector_fill: str | tuple[int, ...] = "gray",
     connector_width: int = 1,
     ellipsize_after: int = 10,
-    inout_as_tensor: bool = True,
     show_neurons: bool = True,
+    opacity: int = 255,
 ) -> Image.Image:
-    """
-    Generates an architecture visualization for a given linear PyTorch model (i.e., one input and output tensor for each
-    layer) in a graph style.
+    """Generates an architecture visualization for a given linear PyTorch model in a graph style.
 
     Args:
         model (torch.nn.Module): A PyTorch model that will be visualized.
@@ -39,8 +40,8 @@ def graph_view(
         to_file (str, optional): Path to the file to write the created image to. If the image does not exist yet,
             it will be created, else overwritten. Image type is inferred from the file ending. Providing None
             will disable writing.
-        color_map (dict, optional): Dict defining fill and outline for each layer by class type. Will fallback to default
-            values for not specified classes.
+        color_map (dict, optional): Dict defining fill and outline for each layer by class type. Will fallback
+            to default values for not specified classes.
         node_size (int, optional): Size in pixels each node will have.
         background_fill (Any, optional): Color for the image background. Can be str or (R,G,B,A).
         padding (int, optional): Distance in pixels before the first and after the last layer.
@@ -50,108 +51,59 @@ def graph_view(
         connector_width (int, optional): Line-width of the connectors in pixels.
         ellipsize_after (int, optional): Maximum number of neurons per layer to draw. If a layer is exceeding this,
             the remaining neurons will be drawn as ellipses.
-        inout_as_tensor (bool, optional): If True there will be one input and output node for each tensor, else the
-            tensor will be flattened and one node for each scalar will be created (e.g., a (10, 10) shape will be
-            represented by 100 nodes).
         show_neurons (bool, optional): If True a node for each neuron in supported layers is created (constrained by
             ellipsize_after), else each layer is represented by a node.
+        opacity (int, optional): Transparency of the color (0 ~ 255).
 
     Returns:
         Image.Image: Generated architecture image.
     """
-
-    if color_map is None:
-        color_map = dict()
+    _color_map: dict = {}
+    if color_map is not None:
+        _color_map = defaultdict(dict, color_map)
 
     # Iterate over the model to compute bounds and generate boxes
 
-    layers: List[Any] = list()
-    layer_y = list()
-
     # Attach helper layers
 
     id_to_num_mapping, adj_matrix, model_layers = model_to_adj_matrix(
-        model, input_shape
+        model,
+        input_shape,
     )
 
     # Add fake input layers
 
     id_to_num_mapping, adj_matrix, model_layers = add_input_dummy_layer(
-        input_shape, id_to_num_mapping, adj_matrix, model_layers
+        input_shape,
+        id_to_num_mapping,
+        adj_matrix,
+        model_layers,
     )
 
     # Create architecture
 
     current_x = padding  # + input_label_size[0] + text_padding
 
-    id_to_node_list_map = dict()
-
-    for index, layer_list in enumerate(model_layers):
-        current_y = 0
-        nodes = []
-        layer: Any
-        for layer in layer_list:
-            is_box = True
-            units = 1
-
-            if show_neurons:
-                if hasattr(layer, "_saved_bias_sym_sizes_opt"):
-                    is_box = False
-                    units = layer._saved_bias_sym_sizes_opt[0]
-                elif hasattr(layer, "_saved_mat2_sym_sizes"):
-                    is_box = False
-                    units = layer._saved_mat2_sym_sizes[1]
-                elif hasattr(layer, "units"):  # for dummy input layer
-                    is_box = False
-                    units = layer.units
-
-            n = min(units, ellipsize_after)
-            layer_nodes = list()
-
-            for i in range(n):
-                scale = 1
-                c: Box | Circle | Ellipses
-                if not is_box:
-                    if i != ellipsize_after - 2:
-                        c = Circle()
-                    else:
-                        c = Ellipses()
-                else:
-                    c = Box()
-                    scale = 3
-
-                c.x1 = current_x
-                c.y1 = current_y
-                c.x2 = c.x1 + node_size
-                c.y2 = c.y1 + node_size * scale
-
-                current_y = c.y2 + node_spacing
-
-                c.fill = color_map.get(TARGET_OPS[layer.name()], {}).get(
-                    "fill", "#ADD8E6"
-                )
-                c.outline = color_map.get(TARGET_OPS[layer.name()], {}).get(
-                    "outline", "black"
-                )
-
-                layer_nodes.append(c)
-
-            id_to_node_list_map[str(id(layer))] = layer_nodes
-            nodes.extend(layer_nodes)
-            current_y += 2 * node_size
-
-        layer_y.append(current_y - node_spacing - 2 * node_size)
-        layers.append(nodes)
-        current_x += node_size + layer_spacing
+    layers, layer_y, id_to_node_list_map = _create_architecture(
+        model_layers,
+        current_x,
+        show_neurons,
+        ellipsize_after,
+        node_size,
+        node_spacing,
+        _color_map,
+        opacity,
+        layer_spacing,
+    )
 
     # Generate image
 
-    img_width = (
-        len(layers) * node_size + (len(layers) - 1) * layer_spacing + 2 * padding
-    )
+    img_width = len(layers) * node_size + (len(layers) - 1) * layer_spacing + 2 * padding
     img_height = max(*layer_y) + 2 * padding
     img = Image.new(
-        "RGBA", (int(ceil(img_width)), int(ceil(img_height))), background_fill
+        "RGBA",
+        (int(ceil(img_width)), int(ceil(img_height))),
+        background_fill,
     )
 
     draw = aggdraw.Draw(img)
@@ -164,7 +116,7 @@ def graph_view(
             node.y1 += y_off
             node.y2 += y_off
 
-    for start_idx, end_idx in zip(*np.where(adj_matrix > 0)):
+    for start_idx, end_idx in zip(*np.where(adj_matrix > 0), strict=False):
         start_id = next(get_keys_by_value(id_to_num_mapping, start_idx))
         end_id = next(get_keys_by_value(id_to_num_mapping, end_idx))
 
@@ -172,10 +124,11 @@ def graph_view(
         end_layer_list = id_to_node_list_map[end_id]
 
         # draw connectors
-        for start_node_idx, start_node in enumerate(start_layer_list):
+        for start_node in start_layer_list:
             for end_node in end_layer_list:
                 if not isinstance(start_node, Ellipses) and not isinstance(
-                    end_node, Ellipses
+                    end_node,
+                    Ellipses,
                 ):
                     _draw_connector(
                         draw,
@@ -185,8 +138,8 @@ def graph_view(
                         width=connector_width,
                     )
 
-    for i, layer in enumerate(layers):
-        for node_index, node in enumerate(layer):
+    for layer in layers:
+        for node in layer:
             node.draw(draw)
 
     draw.flush()
@@ -197,10 +150,97 @@ def graph_view(
     return img
 
 
-def _draw_connector(draw, start_node, end_node, color, width):
+def _draw_connector(
+    draw: aggdraw.Draw,
+    start_node: Box | Circle | Ellipses,
+    end_node: Box | Circle | Ellipses,
+    color: str | tuple[int, ...],
+    width: int,
+) -> None:
+    """Draw the line connector between nodes."""
     pen = aggdraw.Pen(color, width)
     x1 = start_node.x2
     y1 = start_node.y1 + (start_node.y2 - start_node.y1) / 2
     x2 = end_node.x1
     y2 = end_node.y1 + (end_node.y2 - end_node.y1) / 2
     draw.line([x1, y1, x2, y2], pen)
+
+
+def _retrieve_isbox_units(layer: torch.autograd.Function, show_neurons: bool) -> tuple[bool, int]:
+    """Return the number of units and the flag whether to visualize using a box or not."""
+    print("test: ", type(layer))
+    is_box = True
+    units = 1
+    if show_neurons:
+        if hasattr(layer, "_saved_bias_sym_sizes_opt"):
+            is_box = False
+            units = layer._saved_bias_sym_sizes_opt[0]  # noqa: SLF001
+        elif hasattr(layer, "_saved_mat2_sym_sizes"):
+            is_box = False
+            units = layer._saved_mat2_sym_sizes[1]  # noqa: SLF001
+        elif hasattr(layer, "units"):  # for dummy input layer
+            is_box = False
+            units = layer.units
+    return is_box, units
+
+
+def _create_architecture(
+    model_layers: list[list],
+    current_x: int,
+    show_neurons: bool,
+    ellipsize_after: int,
+    node_size: int,
+    node_spacing: int,
+    color_map: dict[Any, Any],
+    opacity: int,
+    layer_spacing: int,
+) -> tuple[list, list, dict]:
+    """Create nodes of architecture for each layers."""
+    id_to_node_list_map = {}
+    layers = []
+    layer_y = []
+    for layer_list in model_layers:
+        current_y = 0
+        nodes = []
+        layer: Any
+        for layer in layer_list:
+            is_box, units = _retrieve_isbox_units(layer, show_neurons)
+
+            n = min(units, ellipsize_after)
+            layer_nodes = []
+
+            for i in range(n):
+                scale = 1
+                c: Box | Circle | Ellipses
+                if not is_box:
+                    c = Circle() if i != ellipsize_after - 2 else Ellipses()
+                else:
+                    c = Box()
+                    scale = 3
+
+                c.x1 = current_x
+                c.y1 = current_y
+                c.x2 = c.x1 + node_size
+                c.y2 = c.y1 + node_size * scale
+
+                current_y = c.y2 + node_spacing
+
+                c.set_fill(
+                    color_map.get(TARGET_OPS[layer.name()], {}).get("fill", "#ADD8E6"),
+                    opacity,
+                )
+                c.outline = color_map.get(TARGET_OPS[layer.name()], {}).get(
+                    "outline",
+                    "black",
+                )
+
+                layer_nodes.append(c)
+
+            id_to_node_list_map[str(id(layer))] = layer_nodes
+            nodes.extend(layer_nodes)
+            current_y += 2 * node_size
+
+        layer_y.append(current_y - node_spacing - 2 * node_size)
+        layers.append(nodes)
+        current_x += node_size + layer_spacing
+    return layers, layer_y, id_to_node_list_map