Merge pull request #5 from LucaBonfiglioli/develop

Version 0.4.0
LucaBonfiglioli · Mar 12, 2023 · fa6516e · fa6516e
2 parents 1057e03 + 662af0e
commit fa6516e
Show file tree

Hide file tree

Showing 44 changed files with 1,239 additions and 349 deletions.
diff --git a/.gitignore b/.gitignore
@@ -127,10 +127,14 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# NNViz outputs
 *.pdf
 *.svg
+*.json
 !examples/*.pdf
 !examples/*.svg
+!examples/*.json
 
 # VSCode
 .vscode/

diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md
@@ -0,0 +1,7 @@
+# No Code of Conduct
+
+Everybody is welcome to contribute to this project regardless of their background, ethnicity, sexual orientation, religion or political views or any other set of characteristics that make them unique.
+
+I do care instead about is building something that is useful, nothing else matters. I expect you to be respectful and to be able to discuss your ideas in a constructive way, but remember that this is not a safe space and that I will not enforce any additional set of rules beyond the ones that already in place.
+
+Please refrain from contributing to this project if you feel that you can be easily offended by the content of the code or the discussion, because nothing will be done to protect you from that unless I am forced to do so.
diff --git a/README.md b/README.md
diff --git a/examples/resnet18_1.pdf b/examples/resnet18_1.pdf
diff --git a/examples/resnet18_1_edge.pdf b/examples/resnet18_1_edge.pdf
diff --git a/examples/resnet18_2.pdf b/examples/resnet18_2.pdf
diff --git a/examples/resnet18_2_edge.pdf b/examples/resnet18_2_edge.pdf
diff --git a/examples/resnet18_full.pdf b/examples/resnet18_full.pdf
diff --git a/examples/resnet18_full_edge.pdf b/examples/resnet18_full_edge.pdf
diff --git a/examples/stupid_conv_full.pdf b/examples/stupid_conv_full.pdf
diff --git a/examples/stupid_conv_full_edge.pdf b/examples/stupid_conv_full_edge.pdf
diff --git a/examples/stupid_model.py b/examples/stupid_model.py
@@ -1,8 +1,8 @@
+import typing as t
+
 import torch
 import torch.nn as nn
 
-import typing as t
-
 
 class NormalLayer(nn.Module):
     def __init__(self, in_features: int, out_features: int) -> None:
@@ -77,10 +77,12 @@ def forward(
 
 # Generate stupid model graph with:
 
-# nnviz examples\stupid_model.py:StupidModel -d 1 -s -o examples\stupid_model_1_edge.pdf \
+# nnviz examples\stupid_model.py:StupidModel \
+# -d 1 -s -o examples\stupid_model_1_edge.pdf \
 # -i "{'x': [torch.randn(3, 10), torch.randn(3, 10)], 'y': torch.rand(3, 10)}"
 
 # This is just an example of how cursed things can get when dealing with torch.fx
-# Don't ever write code like this. If you do, WIDELUCA will personally seek you out and make you regret it.
+# Don't ever write code like this. If you do, WIDELUCA will personally seek you out and
+# make you regret it.
 
 # Repent, before it's too late.
diff --git a/examples/stupid_model_1.pdf b/examples/stupid_model_1.pdf
diff --git a/examples/stupid_model_1_edge.pdf b/examples/stupid_model_1_edge.pdf
diff --git a/images/convnext_tiny.jpeg b/images/convnext_tiny.jpeg
diff --git a/images/convnext_tiny_features-3-0-block.jpeg b/images/convnext_tiny_features-3-0-block.jpeg
diff --git a/images/efficientnet_b0.jpeg b/images/efficientnet_b0.jpeg
diff --git a/images/efficientnet_b0_1.jpeg b/images/efficientnet_b0_1.jpeg
diff --git a/images/efficientnet_b0_2.jpeg b/images/efficientnet_b0_2.jpeg
diff --git a/images/efficientnet_b0_3.jpeg b/images/efficientnet_b0_3.jpeg
diff --git a/images/efficientnet_b0_full.jpeg b/images/efficientnet_b0_full.jpeg
diff --git a/images/list_of_tensors.jpeg b/images/list_of_tensors.jpeg
diff --git a/images/resnet18.jpeg b/images/resnet18.jpeg
diff --git a/images/resnet18_edges.jpeg b/images/resnet18_edges.jpeg
diff --git a/images/resnet18_layer2-0.jpeg b/images/resnet18_layer2-0.jpeg
diff --git a/nnviz/__init__.py b/nnviz/__init__.py
@@ -1 +1 @@
-__version__ = "0.3.1"
+__version__ = "0.4.0"
diff --git a/nnviz/cli/main.py b/nnviz/cli/main.py
@@ -11,34 +11,89 @@
 - A string in the form `PATH_TO_FILE:NAME`, where `PATH_TO_FILE` is the path to a python
 file containing a model, and `NAME` is the name of the model in that file. Valid names
 include objects of type `torch.nn.Module` and functions that return an object of type
-`torch.nn.Module` with no arguments. Class constructors are considered like functions.
+`torch.nn.Module` with no arguments. Class constructors are considered like functions.\n
+- A PATH to a json file containing a previously serialized graph. In this case NNViz
+will simply load the graph (without inspecting anything) and draw it. To create such a
+file, use the `-j` or `--json` flag when running this command.\n
 """
-out_help = "The output file path. If not provided, it will save a pdf file named after the model in the current directory."
-depth_help = "The maximum depth of the graph. No limit if < 0."
+out_help = """
+The output file path. If not provided, it will save a pdf file named after the model in
+the current directory.
+"""
+depth_help = "The maximum depth of the graph. No limit if < 0. Default is 2."
 show_help = "Also show the graph after drawing using the default pdf application."
-input_help = """The input to feed to the model. If specified, nnviz will also add synthetic
-representation of the data passing through the model. Can either be: \n
+input_help = """The input to feed to the model. If specified, nnviz will also add
+synthetic representation of the data passing through the model. Can either be: \n
 - "default" -> float32 BCHW tensor of shape (1, 3, 224, 224) (commonly used) \n
 - "image<side>" (e.g. image224, image256, ...) -> float32 BCHH tensor \n
-- "image<height>x<width>" (e.g. image224x224, image256x512, ...) -> float32 BCHW tensor \n
+- "image<height>x<width>" (e.g. image224x224, image32x64, ...) -> float32 BCHW tensor\n
 - "tensor<s0>x<s1>x<s2>x..." (e.g. tensor1x3x224x224, tensor1x3x256x512, ...) ->
 float32 generic tensors \n
 - "<key1>:<value1>;<key2>:<value2>;... (e.g. x:tensor1x3x224x224;y:tensor1x3x256x512,
 ...) -> dictionary of tensors \n
-- A plain python string that evaluates to a dictionary of tensors (e.g. "{'x': torch.rand(1, 3, 224, 224)}")
+- A plain python string that evaluates to a dictionary of tensors
+(e.g. "{'x': torch.rand(1, 3, 224, 224)}")
+"""
+layer_help = """
+The name of the layer to visualize. If not provided, the whole model
+will be visualized.
+"""
+json_help = "Also save the graph as a json file, named just like the output file."
+collapse_help = """
+Layers that should collapsed, besides the ones that are collapsed by depth.
+"""
+quiet_help = "Disable all printing besides eventual errors."
+style_help = """
+List of style options to apply, in the form `key=value`. Don't worry about the type of
+the value, it will be automatically inferred by pydantic. The available options are: \n
+- fontname: str - The font to use for the graph. Default is "Arial". \n
+- default_node_color: str - The default color for nodes (in case the colorizer fails to
+return a color). Default is "gray". \n
+- default_edge_color: str - The default color for edges. Default is "black". \n
+- node_style: str - The style for nodes. See graphviz docs for details. Default is
+"rounded,filled". \n
+- node_margin: str - The horizontal and vertical margin for nodes. See graphviz docs for
+details. Default is "0.2,0.1". \n
+- edge_thickness: str - The thickness of edges. Default is "2.0". \n
+- graph_title_font_size: int - The font size for the graph title. Default is 48. \n
+- node_title_font_size: int - The font size for the node title. Default is 24. \n
+- cluster_title_font_size: int - The font size for the cluster title. Default is 18. \n
+- show_title: bool - Whether to show the graph title. Default is True. \n
+- show_specs: bool - Whether to show the specs as a label for each edge. 
+Default is True. \n
+- show_node_name: bool - Whether to show the node name (just below the title).  
+Default is True. \n
+- show_node_params: bool - Whether to show the count of parameters for each node. 
+Default is True. \n
+- show_node_arguments: bool - Whether to show the arguments for each node. 
+Default is True. \n
+- show_node_source: bool - Whether to show the source of each node. Default is True. \n
+- show_clusters: bool - Whether to show the clusters as gray subgraphs. 
+Default is True. \n
 """
 
 
 @app.command(name="quick")
 def quick(
+    # Arguments
     model: str = typer.Argument(..., help=model_help),
-    output_path: t.Optional[Path] = typer.Option(None, "-o", "--out", help=out_help),
-    depth: int = typer.Option(1, "-d", "--depth", help=depth_help),
-    show: bool = typer.Option(False, "-s", "--show", help=show_help),
+    # Options
+    layer: t.Optional[str] = typer.Option(None, "-l", "--layer", help=layer_help),
     input: str = typer.Option(None, "-i", "--input", help=input_help),
+    depth: int = typer.Option(2, "-d", "--depth", help=depth_help),
+    collapse: t.List[str] = typer.Option([], "-c", "--collapse", help=collapse_help),
+    output: t.Optional[Path] = typer.Option(None, "-o", "--out", help=out_help),
+    # Style
+    style: t.List[str] = typer.Option([], "-S", "--style", help=style_help),
+    # Flags
+    show: bool = typer.Option(False, "-s", "--show", help=show_help),
+    json: bool = typer.Option(False, "-j", "--json", help=json_help),
+    quiet: bool = typer.Option(False, "-q", "--quiet", help=quiet_help),
 ) -> None:
     """Quickly visualize a model."""
-    from nnviz import drawing, inspection
+    from traceback import print_exc
+
+    from nnviz import drawing, entities, inspection
 
     def _show(output_path: Path) -> None:
         import os
@@ -58,37 +113,123 @@ def _show(output_path: Path) -> None:
             subprocess.Popen(["xdg-open", output_path])
 
         else:
-            typer.echo(f"Could not open {output_path} automatically. I'm sorry.")
+            RuntimeError("Could not automatically open the file.")
 
-    if output_path is None:
+    if output is None:
         _, _, model_name = model.rpartition(":")
-        output_path = Path(f"{model_name}.pdf")
-
-    # Load model
+        if layer is not None:
+            model_name += f"_{layer}".replace(".", "-")
+        output = Path(model_name).with_suffix(".pdf")
+
+    # If the model is a json file, load it directly without inspecting
+    if model.endswith(".json"):
+        if not quiet:
+            typer.echo(f"Loading graph from json file ({model})...")
+        try:
+            graph_data = entities.GraphData.parse_file(model)
+            graph = entities.NNGraph.from_data(graph_data)
+        except Exception as e:
+            print_exc()
+            raise typer.BadParameter(
+                f"Could not load graph from json file {model}"
+            ) from e
+
+        # Remove .json from model name (mildly cursed)
+        model = model[:-5]
+
+    # Otherwise, inspect the model
+    else:
+        # Load model
+        if not quiet:
+            typer.echo(f"Parsing model '{model}'...")
+        try:
+            nn_model = inspection.load_from_string(model)
+        except Exception as e:
+            print_exc()
+            raise typer.BadParameter(f"Could not load model {model}") from e
+
+        # Get layer if needed
+        if layer is not None:
+            try:
+                nn_model = nn_model.get_submodule(layer)
+            except Exception as e:
+                raise typer.BadParameter(f"Could not find layer {layer}") from e
+
+        # Create inspector
+        inspector = inspection.TorchFxInspector()
+
+        # Parse input
+        if not quiet and input is not None:
+            typer.echo(f"Parsing input '{input}'...")
+        try:
+            parsed_input = inspection.parse_input_str(input)
+        except Exception as e:
+            print_exc()
+            raise typer.BadParameter(f"Could not parse input {input}") from e
+
+        # Inspect
+        if not quiet:
+            typer.echo("Inspecting model...")
+        try:
+            graph = inspector.inspect(nn_model, inputs=parsed_input)
+        except Exception as e:
+            print_exc()
+            raise typer.BadParameter(f"Could not inspect model {model}") from e
+
+    # Save json if needed
+    if json:
+        json_path = output.with_suffix(".json")
+        if not quiet:
+            typer.echo(f"Saving graph as json file ({json_path})...")
+        try:
+            with open(json_path, "w") as f:
+                f.write(graph.data.json())
+        except Exception as e:
+            print_exc()
+            raise typer.BadParameter(
+                f"Could not save graph as json file {json_path}"
+            ) from e
+
+    if not quiet and (depth > 0 or len(collapse) > 0):
+        typer.echo("Collapsing graph...")
     try:
-        nn_model = inspection.load_from_string(model)
-    except Exception as e:
-        raise typer.BadParameter(f"Could not load model {model}") from e
-
-    # TODO: The inspector and drawer should be configurable, not hardcoded
-    inspector = inspection.TorchFxInspector()
-    drawer = drawing.GraphvizDrawer(output_path)
-
-    # Parse input
-    parsed_input = inspection.parse_input_str(input)
-
-    # Inspect
-    graph = inspector.inspect(nn_model, inputs=parsed_input)
+        # Collapse by depth
+        graph.collapse_by_depth(depth)
 
-    # Collapse by depth
-    graph = graph.collapse(depth)
+        # Collapse by name
+        graph.collapse_multiple(collapse)
+    except Exception as e:
+        print_exc()
+        raise typer.BadParameter("Could not collapse graph") from e
 
     # Draw
-    drawer.draw(graph)
+    if not quiet:
+        typer.echo(f"Drawing graph to {output}...")
+    try:
+        parsed_style_dict = {}
+        for s in style:
+            key, _, value = s.partition("=")
+            parsed_style_dict[key] = value
+        parsed_style = drawing.GraphvizDrawerStyle.parse_obj(parsed_style_dict)
+        drawer = drawing.GraphvizDrawer(output, style=parsed_style)
+        drawer.draw(graph)
+    except Exception as e:
+        print_exc()
+        raise typer.BadParameter(f"Could not draw graph to {output}") from e
 
     # Show
     if show:
-        _show(output_path)
+        if not quiet:
+            typer.echo(f"Opening {output}...")
+        try:
+            _show(output)
+        except Exception as e:
+            print_exc()
+            raise typer.BadParameter(f"Could not open {output}") from e
+
+    if not quiet:
+        typer.echo()
+        typer.echo("Done!")
 
 
 if __name__ == "__main__":

diff --git a/nnviz/colors/__init__.py b/nnviz/colors/__init__.py
@@ -1,3 +1,3 @@
-from nnviz.colors.base import RGBColor, ColorPicker
-from nnviz.colors.hashcolor import HashColorPicker
+from nnviz.colors.base import ColorPicker, RGBColor
 from nnviz.colors.bubble import BubbleColorPicker
+from nnviz.colors.hashcolor import HashColorPicker
diff --git a/nnviz/colors/bubble.py b/nnviz/colors/bubble.py
@@ -117,8 +117,8 @@ def spawn(self, key: t.Hashable) -> None:
             # Select the closest child
             min_index, min_distance = np.argmin(distances), np.min(distances)
 
-            # If the distance to the edge is smaller than the half distance to the closest
-            # child, the direction is towards the center of the parent tree
+            # If the distance to the edge is smaller than the half distance to the
+            # closest child, the direction is towards the center of the parent tree
             if distance_to_edge < min_distance:
                 direction = self.origin - new_origin
 
@@ -194,7 +194,6 @@ def pick(self, *args: t.Hashable) -> colors.RGBColor:
     possible_keys = [chr(i) for i in range(ord("a"), ord("a") + 21)]
 
     while True:
-
         canvas = np.zeros((H, W, 3))
 
         # Sample with replacement from the possible keys