Package: nnviz
📊 NNVizNNViz is a python package to visualize neural networks in a comprehensible way. It is based onGraphvizandPyTorch.More PDFs examples can be found in theexamplesfolder.🧰 InstallationMake sure you have a working installation ofGraphvizon your system.If you are using Windows, graphviz dependencies are easily installable with conda, so consider usingMinicondaorAnacondainstead of the plain Python distribution. Trying to install graphviz without conda is generally associated with symptoms like headaches, asphyxiation, and death.Linux users can install graphviz (and graphviz-dev) with their package manager (e.g.sudo apt install -y graphviz graphviz-devon Ubuntu).Then:pipinstallnnviz🖥️ CLI UsageCurrently, nnviz is supposed to be used via its CLI, and not as a library. A proper API will be added in the future, most likely with the 1.x release. Anyway, the CLI can do a lot of things, and it is pretty easy to use.I will try to explain the most important things here, but you can always display everything there is to know by entering:nnviz--helpThe only entry point of NNViz is thennvizcommand that in its simplest form does roughly the following:Load a model from torchvision, a local file or an installed package.Perform a symbolic trace of the model and extract a graph representation of it.Draw the graph representation to a pdf file named after the model.Additional behavior is controlled by a set of options and flags that can be passed to the command.🔥 Visualizing built-in torchvision modelsYou can visualize any model from torchvision by passing its name to thennvizcommand.nnvizresnet18
nnvizefficientnet_b0
nnvizconvnext_tinyThese models are saved to a pdf file named after the model in the current working directory.📦 Visualizing models from installed packagesBeing limited to torchvision models is not very useful, so you can also visualize models from any installed package, as long as they are importable. To do so, you need to pass the fully qualified name of a python function (or class constructor) that returns an instance of the model you want to visualize.nnvizpackage.module.submodule:MyModelConstructor
nnvizpackage.module.submodule:a_function_that_returns_a_model⚠️NOTE:Currently, there is a limitation with this feature: the model constructormusttake no arguments (or at least no required ones), because I cannot find a way to easily pass arguments to the constructor from the command line. A workaround is to manually create a python file with a function that returns the model, and then pass that function to thennvizcommand. I know, it's not ideal, this will be fixed in the 0.5.x release.🐍 Visualizing models from python filesAnd if want to visualize a model that is not in a package, you can also pass a path to a python file that contains the model constructor. NNViz will attempt to dynamically import the file and extract the model constructor from it.nnvizpath/to/file.py:MyModelConstructor
nnvizpath/to/file.py:a_function_that_returns_a_model⚠️NOTE:Currently, there is a limitation with this feature: the model constructormusttake no arguments (or at least no required ones), because I cannot find a way to easily pass arguments to the constructor from the command line. A workaround is to manually create a python file with a function that returns the model, and then pass that function to thennvizcommand. I know, it's not ideal, this will be fixed in the 0.5.x release.📑 Choosing a different output fileSaving the pdf to the current working directory is not very useful, so you can also pass a path to a different file to thennvizcommand using the-oor--outoption.nnvizresnet18-omy_resnet18.pdf
nnvizresnet18-omy_resnet18.png
nnvizresnet18-omy_resnet18.svg
nnvizresnet18--out/very/custom/path/to/wow.pdf📜 Dumping/Loading the graph representationIf you plan to create multiple visualizations of the same model, you can use JSON files to store the graph representation of the model, and then load it from the file instead of performing a symbolic trace every time, which, if done repeatedly, can become quite slow.To dump the graph representation to a JSON file, you can pass the-jor--jsonflag to thennvizcommand.
The JSON file will be saved together with the pdf file, just with a ".json" extension instead of ".pdf".nnvizresnet18-j
nnvizresnet18--json
nnvizresnet18-j-o/very/custom/path/to/wow.pdfIn the third example (above), the JSON file will be saved to/very/custom/path/to/wow.json.Ok, but once you have the JSON file, how do you load it? Well, you can pass the path to the JSON file to thennvizcommand as if it was a model name, and it will load the graph representation from the file instead of performing a trace, and it will be blazing fast.nnviz/very/custom/path/to/wow.json🖼️ Displaying the output file (for the laziest people)If you happen to be too lazy to open the pdf file manually, or press theF5key to refresh the file explorer, you can also pass the-sor--showflag to thennvizcommand to automatically open the generated file with the default app associated to the specified file type.nnvizresnet18-s
nnvizresnet18--show🪓 Plotting only a specific submoduleSometimes you may be interested in visualizing only a specific submodule instead of the whole thing. To do so, you can pass the-lor--layeroption to thennvizcommand, followed by the name of the submodule you want to visualize. The submodule can be specified using a dot-separated path, just like you would do in python.nnvizresnet18-llayer2.0
nnvizconvnext_tiny-lfeatures.3.0.blockThis will only plot the specified submodule and all of its children and save the result to a pdf file named after the submodule. For example, in the first example above, the pdf file will be namedresnet18_layer2-0.pdfand in the second,convnext_tiny_features-3-0-block.pdf. You can always change the name of the output file using the-oor--outoption.🌂 Reducing the complexity of the graphWith NNViz, you can also collapse groups of nodes into a single "collapsed" node to reduce the complexity of the graph and improve readability. Currently the CLI allows to do so by two different methods:Collapse by depth: merge together all nodes that have a depth greater than a specified value. The "depth" of a node is the number of nested submodules it is contained in.Collapse by name: merge together all nodes that share the same path prefix. This is useful in case you want a detailed view of a specific submodule, but you don't care about the details of the other ones.The depth collpsing is enabled by default and set to 2, which in most cases is enough to reduce the complexity of the graph without losing too much information. You can change the depth by passing the-dor--depthoption to thennvizcommand, followed by the desired depth.Set the depth to -1 to disable depth collapsing and plot the whole graph.nnvizefficientnet_b0-d1nnvizefficientnet_b0-d2nnvizefficientnet_b0-d3nnvizefficientnet_b0--depth-1To collapse by name, you can pass the-cor--collapseoption to thennvizcommand, followed by a path prefix. The path prefix can be specified using a dot-separated path, just like you would do in python. You can use this option multiple times to collapse multiple groups of nodes.nnvizefficientnet_b0--depth-1--collapsefeaturesHere, the "classifier" will be fully displayed, but the "features" submodule will be collapsed into a single node.🧮 Data tracingBesides the graph representation of the model, NNViz can also trace the data flow through the model and represent the data as labels attached to each graph edge. This is useful to understand how the data is transformed by the model and to identify potential bottlenecks in the model.Data tracing cannot be done with a simple symbolic trace, because the data is not available at that point. We need to actually run the model on some bogus data to get the data flow, which can be more expensive in terms of both time and resources. For this reason data tracing is disabled by default.To enable data tracing, just provide an input to the model. You can do with the-ior--inputoption to thennvizcommand, followed by a string representing the input you want to pass.Parsing complex tensor data from CLI is not easy, NNViz will try its best, but you may need to pass the input data as python code.Supported input types:default: The default tensor used most commonly for image classification models. Equivalent to python code:{"x":torch.randn(1,3,224,224)}image<side>: A square image with the specified side (e.g.image256,image64,image241...). Equivalent to python code:{"x":torch.randn(1,3,side,side)}image<height>x<width>: An image with the specified height and width (e.g.image256x256,image64x64,image241x241...). Equivalent to python code:{"x":torch.randn(1,3,height,width)}tensor<shape>: A tensor with the specified shape (e.g.tensor18x23,tensor30,tensor4x5x256...). Equivalent to python code:{"x":torch.randn(*shape)}<key1>:<value1>;<key2>:<value2>;...: A dictionary of key-value pairs. Keys are strings and values can be any of the above types. This is especially useful for models that take multiple inputs, or have an input name different fromx.Plain python code. If the input string is not recognized as one of the above types, it will be evaluated as python code. This is useful if you want to pass a more complex input, like a list of tensors, or a dictionary of tensors. When using python code you will have the moduletorchalready imported, so you can use it to create tensors. Example:{"x":torch.randn(1,3,224,224),"y":[torch.randn(1,3,224,224),torch.randn(1,3,224,224)]}⚠️NOTE:You cannot just spam "default" into every model that comes to your mind. To perform data tracing you need to pass something that makes sense for the model. For example, if you pass an image to a linear model, the data tracing will fail, so it will default to the basic graph representation.Let's see some examples:nnvizresnet18-idefaultAs you can see, every edge is now labeled with the shape of the tensor that flows through it.And what if the data is not a tensor? I thought of that too! NNViz can also represent nested lists and dictionaries as HTML tables.🎨 Customizing the graphNNViz also allows you to customize the graph representation of the model using the-Sor--styleoption. Currently there are limited style options, but I plan to add more in the future. You can pass multiple style options at once, providing the-Soption multiple times.Style options are passed as a string in the formkey=value. Don't worry about the type of the value, it will be automatically inferred by pydantic. The available options are:fontname: str - The font to use for the graph. Default is "Arial".default_node_color: str - The default color for nodes (in case the colorizer fails to return a color). Default is "gray".default_edge_color: str - The default color for edges. Default is "black".node_style: str - The style for nodes. See graphviz docs for details. Default is "rounded,filled".node_margin: str - The horizontal and vertical margin for nodes. See graphviz docs for details. Default is "0.2,0.1".edge_thickness: str - The thickness of edges. Default is "2.0".graph_title_font_size: int - The font size for the graph title. Default is 48.node_title_font_size: int - The font size for the node title. Default is 24.cluster_title_font_size: int - The font size for the cluster title. Default is 18.show_title: bool - Whether to show the graph title. Default is True.show_specs: bool - Whether to show the specs as a label for each edge. Default is True.show_node_name: bool - Whether to show the node name (just below the title). Default is True.show_node_params: bool - Whether to show the count of parameters for each node. Default is True.show_node_arguments: bool - Whether to show the arguments for each node. Default is True.show_node_source: bool - Whether to show the source of each node. Default is True.show_clusters: bool - Whether to show the clusters as gray subgraphs. Default is True.Why NNVizDo you want to visualize a torch.nn.Module in a comprehensible way? Here are some options for you:NetronNetronis a web-based viewer for neural network, deep learning and machine learning models. It supportsONNX,Keras,CoreML,Caffe,Caffe2,MXNet,Darknet,TensorFlow,PyTorch,TorchScript,CoreMLandPaddlePaddlemodels.Cool huh? No. It sucks ass:You have to compile your models to ONNX runtime - which can be a pain in the ass and 90% of the times requires you to apply changes to your code.Try to visualize a simple resnet18. You will get an endless clusterfuck of incomprehensible nodes that cannot be collapsed or reduced in any way. This is not useful.TorchvizTorchvizis really just a toy project that just displays the backpropagation graph:Module hierarchy is not tracked/displayed.Cannot collapse nodes.Limited to differentiable models.TorchviewTorchviewseems pretty cool. I did not know of its existence when I started developing nnviz, and in fact was developed in the same period.Well, they do pretty much everything that nnviz does, take a look at their style, and come back here. :D