Graph View

Contents

Graph View#

Graph View module for pytorch model visualization.

visualtorch.graph.graph_view(model, input_shape, input_dtype=None, to_file=None, color_map=None, palette='okabe_ito', node_size=50, background_fill='white', padding=10, layer_spacing=250, node_spacing=10, type_ignore=None, outline_width=1, connector_fill='gray', connector_width=1, ellipsize_after=10, show_neurons=True, opacity=255, show_dimension=False, font=None, font_color='black', level_gap=None, show_input=True, show_arrows=False)#

Generates an architecture visualization for a given linear PyTorch model in a graph style.

Parameters:

  • model (torch.nn.Module) – A PyTorch model that will be visualized.

  • input_shape (tuple) – The shape of the input tensor. For a model whose forward() takes multiple separate input tensors, pass a tuple of per-tensor shapes instead, one per positional argument in order, e.g. ((1, 3, 224, 224), (1, 10)).

  • input_dtype (torch.dtype, optional) – The dtype of the dummy input tensor(s) used to trace the model. Defaults to None for every input, giving a uniformly random float tensor. Needed for a model that starts with an integer/bool-input layer (e.g. nn.Embedding, which rejects a float index tensor): pass torch.long, or - for a model with multiple input tensors of different kinds - a tuple of per-tensor dtypes, e.g. (torch.long, None).

  • 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.

  • palette (str, optional) – Named color palette used as the fallback for any layer type not given an explicit override via color_map. One of “okabe_ito” (default, colorblind-safe), “tol_bright”, “tol_muted”, “tab10”, “grayscale”, “nord”, “dracula”, “gruvbox”, “solarized”, “material”, “catppuccin”.

  • 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.

  • type_ignore (list, optional) – List of layer types in the torch model to ignore during drawing.

  • outline_width (int, optional) – Line width in pixels for the shape borders. Defaults to 1.

  • 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.

  • 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).

  • show_dimension (bool, optional) – If True, print each layer’s output shape below it.

  • font (PIL.ImageFont, optional) – Font used for the shape labels. If None, default font will be used.

  • font_color (str or tuple, optional) – Color for the font if used. Can be a string or a tuple (R, G, B, A).

  • level_gap (int, optional) – Vertical spacing in pixels between stacked skip-connection detour routes. If None, defaults to node_size.

  • show_input (bool, optional) – For a single-input model, whether to draw the synthetic “Input” node. Defaults to True. Set False to hide it - e.g. if you’re overlaying your own custom input illustration instead. Has no effect on a multi-input model, where every input is always shown (omitting any of them would make it ambiguous which arrow belongs to which named input). Ignored (input always kept) when the input feeds more than one consumer, e.g. a residual shortcut, since dropping it would silently discard that edge.

  • show_arrows (bool, optional) – If True, draw a small arrowhead at each connector’s downstream endpoint to indicate data-flow direction.

Returns:

Generated architecture image.

Return type:

Image.Image

Contents