Graph View

Contents

Graph View#

Graph View module for pytorch model visualization.

visualtorch.graph.graph_view(model, input_shape, to_file=None, color_map=None, node_size=50, background_fill='white', padding=10, layer_spacing=250, node_spacing=10, connector_fill='gray', connector_width=1, ellipsize_after=10, show_neurons=True, opacity=255)#

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.

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

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

Returns:

Generated architecture image.

Return type:

Image.Image