LeNet Style View

Contents

LeNet Style View#

LeNet Style View module for pytorch model visualization.

visualtorch.lenet_style.lenet_view(model, input_shape, to_file=None, min_z=1, min_xy=10, max_xy=2000, scale_z=1, scale_xy=1, type_ignore=None, color_map=None, one_dim_orientation='z', background_fill='white', padding=10, spacing=10, draw_funnel=True, shade_step=10, font=None, font_color='black', opacity=255, max_channels=100, offset_z=10, level_gap=None, show_dimension=True, show_input=True)#

Generate a LeNet style architecture visualization for a given torch model.

TODO: remove unnecessary arguments for this LeNet style architecture.

Parameters:

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

  • input_shape (tuple) – The shape of the input tensor (default: (1, 3, 224, 224)). 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)).

  • to_file (str, optional) – Path to the file to write the created image. Overwrite if exist. Image type is inferred from the file extension. Providing None will disable writing.

  • min_z (int, optional) – Minimum size in pixels that a layer will have along the z-axis.

  • min_xy (int, optional) – Minimum size in pixels that a layer will have along the x and y axes.

  • max_channels (int, optional) – Maximum number of channels.

  • max_xy (int, optional) – Maximum size in pixels that a layer will have along the x and y axes.

  • scale_z (float, optional) – Scalar multiplier for the size of each layer along the z-axis.

  • scale_xy (float, optional) – Scalar multiplier for the size of each layer along the x and y axes.

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

  • color_map (dict, optional) – Dictionary defining fill and outline colors for each layer by class type. Will fallback to default values for unspecified classes.

  • one_dim_orientation (str, optional) – Axis on which one-dim layers should be drawn. E.g., ‘x’, ‘y’, or ‘z’.

  • background_fill (str or tuple, optional) – Background color for the image. A string or a tuple (R, G, B, A).

  • padding (int, optional) – Distance in pixels before the first and after the last layer.

  • spacing (int, optional) – Spacing in pixels between two layers.

  • draw_funnel (bool, optional) – If True, a funnel will be drawn between consecutive layers.

  • shade_step (int, optional) – Deviation in lightness for drawing shades (only in volumetric view).

  • font (PIL.ImageFont, optional) – Font that will be used for the legend. 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).

  • opacity (int) – Transparency of the color (0 ~ 255).

  • offset_z (int) – control the offset of overlapping between channels.

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

  • show_dimension (bool, optional) – If True (the default), print each layer’s output shape below it. For a model with parallel branches (e.g. multi-branch merges or multiple input tensors), several boxes can share a column and their labels may overlap - set this to False to drop the labels entirely in that case.

  • show_input (bool, optional) – For a single-input model, whether to draw the synthetic “Input” box. 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.

Returns:

An Image object representing the generated architecture visualization.

Return type:

PIL.Image

Contents