From 762fe0c717533ecbd4fab8ac6757dde00f236f25 Mon Sep 17 00:00:00 2001 From: Willy Fitra Hendria Date: Sun, 25 Feb 2024 14:11:27 +0900 Subject: [PATCH 1/5] add graph view implementation --- visualtorch/__init__.py | 3 +- visualtorch/graph.py | 196 +++++++++++++++++++++++++++++++++++++ visualtorch/layer_utils.py | 154 +++++++++++++++++++++++++++++ visualtorch/utils.py | 43 ++++++++ 4 files changed, 395 insertions(+), 1 deletion(-) create mode 100644 visualtorch/graph.py diff --git a/visualtorch/__init__.py b/visualtorch/__init__.py index 957a9b1..d9e19e3 100644 --- a/visualtorch/__init__.py +++ b/visualtorch/__init__.py @@ -1,3 +1,4 @@ from visualtorch.layered import layered_view +from visualtorch.graph import graph_view -__all__ = ["layered_view"] +__all__ = ["layered_view", "graph_view"] diff --git a/visualtorch/graph.py b/visualtorch/graph.py new file mode 100644 index 0000000..f783a3d --- /dev/null +++ b/visualtorch/graph.py @@ -0,0 +1,196 @@ +import aggdraw +from PIL import Image +from math import ceil +from .layer_utils import model_to_adj_matrix, add_input_dummy_layer +from .utils import Circle, Ellipses, get_keys_by_value, Box +import numpy as np +from typing import Optional, Dict, Any, Tuple, List +import torch + + +def graph_view( + model: torch.nn.Module, + input_shape: Tuple[int, ...], + to_file: Optional[str] = None, + color_map: Optional[Dict[Any, Any]] = None, + node_size: int = 50, + background_fill: Any = "white", + padding: int = 10, + layer_spacing: int = 250, + node_spacing: int = 10, + connector_fill: Any = "gray", + connector_width: int = 1, + ellipsize_after: int = 10, + inout_as_tensor: bool = True, + show_neurons: bool = True, +) -> 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. + + Args: + model (torch.nn.Module): A PyTorch model that will be visualized. + input_shape (tuple): The shape of the input tensor. + 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. + 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. + layer_spacing (int, optional): Spacing in pixels between two layers. + node_spacing (int, optional): Spacing in pixels between nodes. + connector_fill (Any, optional): Color for the connectors. Can be str or (R,G,B,A). + 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. + + Returns: + Image.Image: Generated architecture image. + """ + + if color_map is None: + color_map = dict() + + # 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 + ) + + # 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 + ) + + # 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 = [] + 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(type(layer), {}).get("fill", "orange") + c.outline = color_map.get(type(layer), {}).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 + + # Generate image + + 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 + ) + + draw = aggdraw.Draw(img) + + # y correction (centering) + for i, layer in enumerate(layers): + y_off = (img.height - layer_y[i]) / 2 + node: Any + for node in layer: + node.y1 += y_off + node.y2 += y_off + + for start_idx, end_idx in zip(*np.where(adj_matrix > 0)): + 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)) + + start_layer_list = id_to_node_list_map[start_id] + end_layer_list = id_to_node_list_map[end_id] + + # draw connectors + for start_node_idx, start_node in enumerate(start_layer_list): + for end_node in end_layer_list: + if not isinstance(start_node, Ellipses) and not isinstance( + end_node, Ellipses + ): + _draw_connector( + draw, + start_node, + end_node, + color=connector_fill, + width=connector_width, + ) + + for i, layer in enumerate(layers): + for node_index, node in enumerate(layer): + node.draw(draw) + + draw.flush() + + if to_file is not None: + img.save(to_file) + + return img + + +def _draw_connector(draw, start_node, end_node, color, width): + 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) diff --git a/visualtorch/layer_utils.py b/visualtorch/layer_utils.py index 6992d02..6c7599d 100644 --- a/visualtorch/layer_utils.py +++ b/visualtorch/layer_utils.py @@ -1,7 +1,161 @@ +import numpy as np + +import torch import torch.nn as nn +from .utils import get_keys_by_value + +from typing import Tuple, Dict, List, Any + + +TARGET_OPS = {"AddmmBackward0", "ConvolutionBackward0"} + class SpacingDummyLayer(nn.Module): def __init__(self, spacing: int = 50): super().__init__() self.spacing = spacing + + +class InputDummyLayer: + def __init__(self, name, units=None): + if units: + self.units = units + self.name = name + + +def model_to_adj_matrix( + model, input_shape +) -> Tuple[Dict[str, int], np.ndarray, List[List[torch.nn.Module]]]: + """ + Extract adjacency matrix representation from a pytorch model. + + Args: + model: PyTorch model. + input_shape (tuple): The shape of the input tensor expected by the model, including batch dim. + + Returns: + tuple: A tuple containing: + - id_to_index_adj_mapping (dict): Mapping from node IDs to their corresponding index in the adjacency matrix. + - adjacency_matrix (numpy.ndarray): The adjacency matrix representing connections between model operations/layers. + - model_layers (list): List of model layers organized by their hierarchy. + """ + dummy_input = torch.rand(input_shape) + output_var = model(dummy_input) + + nodes = [] + edges = [] + id_to_ops = {} + + max_level = [0] + max_level_id = [""] + + def add_nodes(fn, source=None, level=0): + assert not torch.is_tensor(fn) + + if str(type(fn).__name__) in TARGET_OPS: + node_id = str(id(fn)) + id_to_ops[node_id] = fn + if node_id not in nodes: + nodes.append(node_id) + level += 1 + if level > max_level[0]: + max_level[0] = level + max_level_id[0] = node_id + + edges.append((node_id, source)) + source = node_id + + # recurse + if hasattr(fn, "next_functions"): + for u in fn.next_functions: + if u[0] is not None: + add_nodes(u[0], source, level) + + def add_base_tensor(var): + if var.grad_fn: + add_nodes(var.grad_fn) + + if var._is_view(): + add_base_tensor(var._base) + + # Extract nodes and edges for the target ops + # Currently only the ones in the TARGET_OPS are supported. + + # handle multiple outputs + if isinstance(output_var, tuple): + for v in output_var: + add_base_tensor(v) + else: + add_base_tensor(output_var) + + # Create adjacency matrix + adjacency_matrix = np.zeros((len(nodes), len(nodes))) + id_to_index_adj_mapping = {node: idx for idx, node in enumerate(nodes)} + + for src_id, trg_id in edges: + if trg_id is not None: + src_index = id_to_index_adj_mapping[src_id] + trg_index = id_to_index_adj_mapping[trg_id] + adjacency_matrix[src_index, trg_index] += 1 + + # Retrieve layers per level + input_layer_id = max_level_id[0] + temp_model_layers = [[input_layer_id]] + + while len(temp_model_layers) < max_level[0]: + prev_layers = temp_model_layers[-1] + new_layer = [] + for layer_id in prev_layers: + src_index = id_to_index_adj_mapping[layer_id] + for trg_idx in np.where(adjacency_matrix[src_index] > 0)[0]: + trg_id = next(get_keys_by_value(id_to_index_adj_mapping, trg_idx)) + new_layer.append(trg_id) + + temp_model_layers.append(list(new_layer)) + + # Filter duplicate layers + seen = set() + model_layers: List[List] = [] + for i in range(len(temp_model_layers) - 1, -1, -1): + new_layers = [] + for layer_id in temp_model_layers[i]: + if layer_id in seen: + continue + seen.add(layer_id) + new_layers.append(id_to_ops[layer_id]) + model_layers.insert(0, list(new_layers)) + + return id_to_index_adj_mapping, adjacency_matrix, model_layers + + +def add_input_dummy_layer( + input_shape: Tuple[int, ...], + id_to_num_mapping: Dict[str, int], + adj_matrix: np.ndarray, + model_layers: List[List[Any]], +) -> Tuple[Dict[str, int], np.ndarray, List[List[str]]]: + """ + Add an input dummy layer to the model layers and update the adjacency matrix accordingly. + + Args: + input_shape (tuple): The shape of the input tensor. + id_to_num_mapping (dict): Mapping from node IDs to their corresponding index in the adjacency matrix. + adj_matrix (numpy.ndarray): The adjacency matrix representing connections between model operations. + model_layers (list): List of model layers organized by their dependencies. + + Returns: + tuple: A tuple containing: + - id_to_num_mapping (dict): Updated mapping from node IDs to their corresponding index in the adjacency matrix. + - adj_matrix (numpy.ndarray): Updated adjacency matrix. + - model_layers (list): Updated list of model layers with the input dummy layer. + """ + first_hidden_layer = model_layers[0][0] + input_dummy_layer = InputDummyLayer("input", input_shape[1]) + model_layers.insert(0, [input_dummy_layer]) + id_to_num_mapping[str(id(input_dummy_layer))] = len(id_to_num_mapping.keys()) + adj_matrix = np.pad( + adj_matrix, ((0, 1), (0, 1)), mode="constant", constant_values=0 + ) + adj_matrix[-1, id_to_num_mapping[str(id(first_hidden_layer))]] += 1 + return id_to_num_mapping, adj_matrix, model_layers diff --git a/visualtorch/utils.py b/visualtorch/utils.py index 329c4ca..18f4b0a 100644 --- a/visualtorch/utils.py +++ b/visualtorch/utils.py @@ -103,6 +103,49 @@ def draw(self, draw: ImageDraw): draw.rectangle([self.x1, self.y1, self.x2, self.y2], pen, brush) +class Circle(RectShape): + def draw(self, draw: ImageDraw): + pen, brush = self._get_pen_brush() + draw.ellipse([self.x1, self.y1, self.x2, self.y2], pen, brush) + + +class Ellipses(RectShape): + def draw(self, draw: ImageDraw): + pen, brush = self._get_pen_brush() + w = self.x2 - self.x1 + d = int(w / 7) + draw.ellipse( + [ + self.x1 + (w - d) / 2, + self.y1 + 1 * d, + self.x1 + (w + d) / 2, + self.y1 + 2 * d, + ], + pen, + brush, + ) + draw.ellipse( + [ + self.x1 + (w - d) / 2, + self.y1 + 3 * d, + self.x1 + (w + d) / 2, + self.y1 + 4 * d, + ], + pen, + brush, + ) + draw.ellipse( + [ + self.x1 + (w - d) / 2, + self.y1 + 5 * d, + self.x1 + (w + d) / 2, + self.y1 + 6 * d, + ], + pen, + brush, + ) + + class ColorWheel: def __init__(self, colors: list | None = None): self._cache: Dict[type, Any] = dict() From 4f409553e3c3ed237e5df91722adb9dd1d246a25 Mon Sep 17 00:00:00 2001 From: Willy Fitra Hendria Date: Sun, 25 Feb 2024 14:22:14 +0900 Subject: [PATCH 2/5] Update description --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 7e62c1a..a3279be 100644 --- a/README.md +++ b/README.md @@ -2,11 +2,11 @@ [![python](https://img.shields.io/badge/python-3.10%2B-blue)]() [![pytorch](https://img.shields.io/badge/pytorch-2.0%2B-orange)]() [![Downloads](https://static.pepy.tech/personalized-badge/visualtorch?period=total&units=international_system&left_color=grey&right_color=green&left_text=PyPI%20Downloads)](https://pepy.tech/project/visualtorch) [![Run Tests](https://github.com/willyfh/visualtorch/actions/workflows/pytest.yml/badge.svg)](https://github.com/willyfh/visualtorch/actions/workflows/pytest.yml) -**VisualTorch** aims to help visualize Torch-based neural network architectures. Currently, this package supports generating layered-style architectures for Torch Sequential and Custom models. This package is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras) by [@paulgavrikov](https://github.com/paulgavrikov). +**VisualTorch** aims to help visualize Torch-based neural network architectures. Currently, this package supports generating layered-style and graph-style architectures for PyTorch Sequential and Custom models. This package is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary). -**v0.2**: Support for custom models has been added. +**v0.2**: Added support for custom models and implemented graph view functionality. -**v0.1.1**: Support for the layered architecture of Torch Sequential. +**v0.1.1**: Added support for the layered architecture of Torch Sequential. ## Installation From 76652c3f73bcd2d3f26c10e553ef01e41ecbbcc9 Mon Sep 17 00:00:00 2001 From: Willy Fitra Hendria Date: Sun, 25 Feb 2024 14:33:16 +0900 Subject: [PATCH 3/5] Add usage example for graph view --- README.md | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index a3279be..f707ffa 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,36 @@ visualtorch.layered_view(model, input_shape=input_shape, legend=True).show() # d ![simple-cnn-custom](https://github.com/willyfh/visualtorch/assets/5786636/f22298b4-f341-4a0d-b85b-11f01e207ad8) +### Graph View +```python +import torch +import torch.nn as nn +import visualtorch + +class SimpleDense(nn.Module): + def __init__(self): + super(SimpleDense, self).__init__() + self.h0 = nn.Linear(4, 8) + self.h1 = nn.Linear(8, 8) + self.h2 = nn.Linear(8, 4) + self.out = nn.Linear(4, 2) + + def forward(self, x): + x = self.h0(x) + x = self.h1(x) + x = self.h2(x) + x = self.out(x) + return x + +model = SimpleDense() + +input_shape = (1, 4) + +visualtorch.graph_view(model, input_shape).show() +``` + +![graph-view](https://github.com/willyfh/visualtorch/assets/5786636/a65b4208-72da-497b-b6c9-aafc82b67b58) + ### Save the Image ```python @@ -142,7 +172,7 @@ Please feel free to send a pull request to contribute to this project. This poject is available as open source under the terms of the [MIT License](https://github.com/willyfh/visualtorch/blob/update-readme/LICENSE). -Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license). +Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary), both of which are also licensed under the MIT license. ## Citation From 9d0336ed2190d1177cc39aa8f51eff9182a49bdb Mon Sep 17 00:00:00 2001 From: Willy Fitra Hendria Date: Sun, 25 Feb 2024 14:35:41 +0900 Subject: [PATCH 4/5] modified by prettier --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index f707ffa..2c3815e 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,7 @@ visualtorch.layered_view(model, input_shape=input_shape, legend=True).show() # d ![simple-cnn-custom](https://github.com/willyfh/visualtorch/assets/5786636/f22298b4-f341-4a0d-b85b-11f01e207ad8) ### Graph View + ```python import torch import torch.nn as nn From 6b2c74c09cf7be81a44c6a5caa5aefee5ede93c1 Mon Sep 17 00:00:00 2001 From: Willy Fitra Hendria Date: Sun, 25 Feb 2024 14:36:42 +0900 Subject: [PATCH 5/5] change graph color --- visualtorch/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/visualtorch/graph.py b/visualtorch/graph.py index f783a3d..e0e3d97 100644 --- a/visualtorch/graph.py +++ b/visualtorch/graph.py @@ -121,7 +121,7 @@ def graph_view( current_y = c.y2 + node_spacing - c.fill = color_map.get(type(layer), {}).get("fill", "orange") + c.fill = color_map.get(type(layer), {}).get("fill", "blue") c.outline = color_map.get(type(layer), {}).get("outline", "black") layer_nodes.append(c)