Package: torchview
torchviewTorchview provides visualization of pytorch models in the form of visual graphs. Visualization includes tensors, modules, torch.functions and info such as input/output shapes.Pytorch version ofplot_model of keras(and more)Supports PyTorch versions $\geq$ 1.7.Useful featuresInstallationFirst, you need to install graphviz,pipinstallgraphvizFor python interface of graphiz to work, you need to have dot layout command working in your system. If it isn't already installed, I suggest you run the following depeding on your OS,Debian-based Linux distro (e.g. Ubuntu):apt-getinstallgraphvizWindows:chocoinstallgraphvizmacOSbrewinstallgraphvizsee more detailshereThen, continue with installing torchview using pippipinstalltorchviewor if you want via condacondainstall-cconda-forgetorchviewor if you want most up-to-date version, install directly from repopipinstallgit+https://github.com/mert-kurttutan/torchview.gitHow To Usefromtorchviewimportdraw_graphmodel=MLP()batch_size=2# device='meta' -> no memory is consumed for visualizationmodel_graph=draw_graph(model,input_size=(batch_size,128),device='meta')model_graph.visual_graphNotebook ExamplesFor more examples, see colab notebooks below,Introduction Notebook:Computer Vision Models:NLP Models:Note:Output graphviz visuals return images with desired sizes. But sometimes, on VSCode, some shapes are being cropped due to large size and svg rendering on by VSCode. To solve this, I suggest you run the followingimportgraphvizgraphviz.set_jupyter_format('png')This problem does not occur on other platforms e.g. JupyterLab or Google Colab.Supported FeaturesAlmost all the models, RNN, Sequentials, Skip Connection, Hugging Face ModelsSupport for meta tensors -> no memory consumption (for very Large models) (pytorch version $\geq$ 1.13) .Shows operations between tensors (in addition to module calls)Rolling/Unrolling feature. Recursively used modules can be rolled visually, see below.Diverse set of inputs/output types, e.g. nested data structure (dict, list, etc), huggingface tokenizer outputsDocumentationdefdraw_graph(model:nn.Module,input_data:INPUT_DATA_TYPE|None=None,input_size:INPUT_SIZE_TYPE|None=None,graph_name:str='model',depth:int|float=3,device:torch.device|str|None=None,dtypes:list[torch.dtype]|None=None,mode:str|None=None,strict:bool=True,expand_nested:bool=False,graph_dir:str|None=None,hide_module_functions:bool=True,hide_inner_tensors:bool=True,roll:bool=False,show_shapes:bool=True,save_graph:bool=False,filename:str|None=None,directory:str='.',**kwargs:Any,)->ComputationGraph:'''Returns visual representation of the input Pytorch Module withComputationGraph object. ComputationGraph object contains:1) Root nodes (usually tensor node for input tensors) which connect to allthe other nodes of computation graph of pytorch module recorded during forwardpropagation.2) graphviz.Digraph object that contains visual representation of computationgraph of pytorch module. This graph visual shows modules/ module hierarchy,torch_functions, shapes and tensors recorded during forward prop, for examplessee documentation, and colab notebooks.Args:model (nn.Module):Pytorch model to represent visually.input_data (data structure containing torch.Tensor):input for forward method of model. Wrap it in a list formultiple args or in a dict or kwargsinput_size (Sequence of Sizes):Shape of input data as a List/Tuple/torch.Size(dtypes must match model input, default is FloatTensors).Default: Nonegraph_name (str):Name for graphviz.Digraph object. Also default name graphviz fileof Graph VisualizationDefault: 'model'depth (int):Upper limit for depth of nodes to be shown in visualization.Depth is measured how far is module/tensor inside the module hierarchy.For instance, main module has depth=0, whereas submodule of main modulehas depth=1, and so on.Default: 3device (str or torch.device):Device to place and input tensors. Defaults togpu if cuda is seen by pytorch, otherwise to cpu.Default: Nonedtypes (list of torch.dtype):Uses dtypes to set the types of input tensor ifinput size is given.mode (str):Mode of model to use for forward prop. Defaultsto Eval mode if not givenDefault: Nonestrict (bool):if true, graphviz visual does not allow multiple edgesbetween nodes. Mutiple edge occurs e.g. when there are tensorsfrom module node to module node and hiding those tensorsDefault: Trueexpand_nested(bool):if true shows nested modules with dashed bordersgraph_dir (str):Sets the direction of visual graph'TB' -> Top to Bottom'LR' -> Left to Right'BT' -> Bottom to Top'RL' -> Right to LeftDefault: None -> TBhide_module_function (bool):Determines whether to hide module torch_functions. Somemodules consist only of torch_functions (no submodule),e.g. nn.Conv2d.True => Dont include module functions in graphvizFalse => Include modules function in graphvizDefault: Truehide_inner_tensors (bool):Inner tensor is all the tensors of computation graphbut input and output tensorsTrue => Does not show inner tensors in graphvizFalse => Shows inner tensors in graphvizDefault: Trueroll (bool):If true, rolls recursive modules.Default: Falseshow_shapes (bool):True => Show shape of tensor, input, and outputFalse => Dont showDefault: Truesave_graph (bool):True => Saves output file of graphviz graphFalse => Does not saveDefault: Falsefilename (str):name of the file to store dot syntax representation andimage file of graphviz graph. Defaults to graph_namedirectory (str):directory in which to store graphviz output files.Default: .Returns:ComputationGraph object that contains visualization of the inputpytorch model in the form of graphviz Digraph object'''ExamplesRolled Version of Recursive Networksfromtorchviewimportdraw_graphmodel_graph=draw_graph(SimpleRNN(),input_size=(2,3),graph_name='RecursiveNet',roll=True)model_graph.visual_graphShow/Hide intermediate (hidden) tensors and Functionals# Show inner tensors and Functionalsmodel_graph=draw_graph(MLP(),input_size=(2,128),graph_name='MLP',hide_inner_tensors=False,hide_module_functions=False,)model_graph.visual_graphResNet / Skip Connection / Support for Torch operations / Nested Modulesimporttorchvisionmodel_graph=draw_graph(resnet18(),input_size=(1,3,32,32),expand_nested=True)model_graph.visual_graphTODODisplay Module parameter infoSupport for Graph Neural Network (GNN)Support for Undirected edges for GNNsSupport for torch-based functions[^1][^1]: Here, torch-based functions refers to any function that uses only torch functions and modules. This is more general than modules.ContributingAll issues and pull requests are much appreciated! If you are wondering how to build the project:torchview is actively developed using the latest version of Python.Changes should be backward compatible to Python 3.7, and will follow Python's End-of-Life guidance for old versions.Runpip install -r requirements-dev.txt. We use the latest versions of all dev packages.To run unit tests, runpytest.To update the expected output files, runpytest --overwrite.To skip output file tests, usepytest --no-outputReferencesParts related to input processing and validation are taken/inspired from torchinfo repository!!.Many of the software related parts (e.g. testing) are also taken/inspired from torchinfo repository. So big thanks to @TylerYep!!!The algorithm for constructing visual graphs is thanks to__torch_function__and subclassingtorch.Tensor. Big thanks to all those who developed this API!!.