Render

Render#

Single public entry point for pytorch model visualization.

graph_view/flow_view/lenet_view (in visualtorch.graph/.flow/.lenet_style) render the same extract_architecture-derived structure three different ways; this module consolidates them behind one function, render(model, input_shape, style=…, **kwargs), so style picks the rendering style and every other parameter is style-appropriate keyword arguments. Kwargs are validated by constructing a per-style dataclass from them - a typo’d kwarg raises TypeError immediately, rather than being silently ignored.

class visualtorch.render.CommonOptions(to_file=None, color_map=None, background_fill='white', padding=10, opacity=255, font=None, font_color='black', level_gap=None)#

Bases: object

Options accepted by every rendering style.

class visualtorch.render.FlowStyleOptions(min_z=10, min_xy=10, max_z=400, max_xy=2000, scale_z=0.1, scale_xy=1, type_ignore=None, one_dim_orientation='z', draw_volume=True, spacing=10, draw_funnel=True, shade_step=10, legend=False, show_dimension=False, show_input=True)#

Bases: object

Options specific to style=”flow” - stacked volumetric/2D boxes connected by funnels.

class visualtorch.render.GraphStyleOptions(node_size=50, layer_spacing=250, node_spacing=10, connector_fill='gray', connector_width=1, ellipsize_after=10, show_neurons=True, show_dimension=False, show_input=True)#

Bases: object

Options specific to style=”graph” - a node/edge diagram, one node per neuron or layer.

class visualtorch.render.LenetStyleOptions(min_z=1, min_xy=10, max_xy=2000, scale_z=1, scale_xy=1, type_ignore=None, one_dim_orientation='z', spacing=10, draw_funnel=True, shade_step=10, max_channels=100, offset_z=10, show_dimension=True, show_input=True)#

Bases: object

Options specific to style=”lenet” - the classic LeNet stacked-plane look.

visualtorch.render.render(model, input_shape, style='graph', **kwargs)#

Generate an architecture visualization for a given PyTorch model.

Parameters:

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

  • input_shape (tuple) – The shape of the input tensor, including batch dim. 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)).

  • style (str, optional) – Which rendering style to use - “graph” (a node/edge diagram), “flow” (stacked volumetric/2D boxes connected by funnels), or “lenet” (the classic LeNet look).

  • **kwargs (Any) – Style-specific and common options (see GraphStyleOptions/FlowStyleOptions/ LenetStyleOptions/CommonOptions for the full list per style). Forwarded into the relevant dataclass constructor, so an unrecognized keyword raises TypeError rather than being silently ignored.

Returns:

Generated architecture image.

Return type:

Image.Image

Contents