analogvnn.utils.render_autograd_graph#

Module Contents#

Classes#

AutoGradDot

Stores and manages Graphviz representation of PyTorch autograd graph.

Functions#

size_to_str(size)

Convert a tensor size to a string.

make_autograd_obj_from_outputs(→ AutoGradDot)

Compile Graphviz representation of PyTorch autograd graph from output tensors.

make_autograd_obj_from_module(→ AutoGradDot)

Compile Graphviz representation of PyTorch autograd graph from forward pass.

get_autograd_dot_from_trace(→ graphviz.Digraph)

Produces graphs of torch.jit.trace outputs.

get_autograd_dot_from_outputs(→ graphviz.Digraph)

Runs and make Graphviz representation of PyTorch autograd graph from output tensors.

get_autograd_dot_from_module(→ graphviz.Digraph)

Runs and make Graphviz representation of PyTorch autograd graph from forward pass.

save_autograd_graph_from_outputs(→ str)

Save Graphviz representation of PyTorch autograd graph from output tensors.

save_autograd_graph_from_module(→ str)

Save Graphviz representation of PyTorch autograd graph from forward pass.

save_autograd_graph_from_trace(→ str)

Save Graphviz representation of PyTorch autograd graph from trace.

analogvnn.utils.render_autograd_graph.size_to_str(size)[source]#

Convert a tensor size to a string.

Parameters:

size (torch.Size) – the size to convert.

Returns:

the string representation of the size.

Return type:

str

class analogvnn.utils.render_autograd_graph.AutoGradDot[source]#

Stores and manages Graphviz representation of PyTorch autograd graph.

Variables:
  • dot (graphviz.Digraph) – Graphviz representation of the autograd graph.

  • _module (nn.Module) – The module to be traced.

  • _inputs (List[Tensor]) – The inputs to the module.

  • _inputs_kwargs (Dict[str, Tensor]) – The keyword arguments to the module.

  • _outputs (Sequence[Tensor]) – The outputs of the module.

  • param_map (Dict[int, str]) – A map from parameter values to their names.

  • _seen (set) – A set of nodes that have already been added to the graph.

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

  • _called (bool) – the module has been called.

property inputs: Sequence[torch.Tensor][source]#

The arg inputs to the module.

Returns:

the arg inputs to the module.

Return type:

Sequence[Tensor]

property inputs_kwargs: Dict[str, torch.Tensor][source]#

The keyword inputs to the module.

Parameters:
  • Dict[str – the keyword inputs to the module.

  • Tensor] – the keyword inputs to the module.

property outputs: Optional[Sequence[torch.Tensor]][source]#

The outputs of the module.

Returns:

the outputs of the module.

Return type:

Optional[Sequence[Tensor]]

property module: torch.nn.Module[source]#

The module.

Returns:

the module to be traced.

Return type:

nn.Module

property ignore_tensor: Dict[int, bool][source]#

The tensor ignored from the dot graphs.

Returns:

the ignore tensor dict.

Return type:

Dict[int, bool]

dot: graphviz.Digraph[source]#
_module: torch.nn.Module[source]#
_inputs: Sequence[torch.Tensor][source]#
_inputs_kwargs: Dict[str, torch.Tensor][source]#
_outputs: Sequence[torch.Tensor][source]#
param_map: dict[source]#
_seen: set[source]#
show_attrs: bool[source]#
show_saved: bool[source]#
max_attr_chars: int[source]#
_called: bool = False[source]#
_ignore_tensor: Dict[int, bool][source]#
__post_init__()[source]#

Create the graphviz graph.

Raises:

ImportError – if graphviz (https://pygraphviz.github.io/) is not available.

reset_params()[source]#

Reset the param_map and _seen.

Returns:

self.

Return type:

AutoGradDot

add_ignore_tensor(tensor: torch.Tensor)[source]#

Add a tensor to the ignore tensor dict.

Parameters:

tensor (Tensor) – the tensor to ignore.

Returns:

self.

Return type:

AutoGradDot

del_ignore_tensor(tensor: torch.Tensor)[source]#

Delete a tensor from the ignore tensor dict.

Parameters:

tensor (Tensor) – the tensor to delete.

Returns:

self.

Return type:

AutoGradDot

get_tensor_name(tensor: torch.Tensor, name: Optional[str] = None) Tuple[str, str][source]#

Get the name of the tensor.

Parameters:
  • tensor (Tensor) – the tensor.

  • name (Optional[str]) – the name of the tensor. Defaults to None.

Returns:

the name and size of the tensor.

Return type:

Tuple[str, str]

add_tensor(tensor: torch.Tensor, name: Optional[str] = None, _attributes=None, **kwargs)[source]#

Add a tensor to the graph.

Parameters:
  • tensor (Tensor) – the tensor.

  • name (Optional[str]) – the name of the tensor. Defaults to None.

  • _attributes (Optional[Dict[str, str]]) – the attributes of the tensor. Defaults to None.

  • **kwargs – the attributes of the dot.node function.

Returns:

self.

Return type:

AutoGradDot

add_fn(fn: Any, _attributes=None, **kwargs)[source]#

Add a function to the graph.

Parameters:
  • fn (Any) – the function.

  • _attributes (Optional[Dict[str, str]]) – the attributes of the function. Defaults to None.

  • **kwargs – the attributes of the dot.node function.

Returns:

self.

Return type:

AutoGradDot

add_edge(u: Any, v: Any, label: Optional[str] = None, _attributes=None, **kwargs)[source]#

Add an edge to the graph.

Parameters:
  • u (Any) – tail node.

  • v (Any) – head node.

  • label (Optional[str]) – the label of the edge. Defaults to None.

  • _attributes (Optional[Dict[str, str]]) – the attributes of the edge. Defaults to None.

  • **kwargs – the attributes of the dot.edge function.

Returns:

self.

Return type:

AutoGradDot

add_seen(item: Any)[source]#

Add an item to the seen set.

Parameters:

item (Any) – the item.

Returns:

self.

Return type:

AutoGradDot

is_seen(item: Any) bool[source]#

Check if the item is in the seen set.

Parameters:

item (Any) – the item.

Returns:

True if the item is in the seen set.

Return type:

bool

analogvnn.utils.render_autograd_graph.make_autograd_obj_from_outputs(outputs: Union[torch.Tensor, Sequence[torch.Tensor]], named_params: Union[Dict[str, Any], Iterator[Tuple[str, torch.nn.Parameter]]], additional_params: Optional[dict] = None, show_attrs: bool = True, show_saved: bool = True, max_attr_chars: int = 50) AutoGradDot[source]#

Compile Graphviz representation of PyTorch autograd graph from output tensors.

Parameters:
  • outputs (Union[Tensor, Sequence[Tensor]]) – output tensor(s) of forward pass

  • named_params (Union[Dict[str, Any], Iterator[Tuple[str, Parameter]]]) – dict of params to label nodes with

  • additional_params (dict) – dict of additional params to label nodes with

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

Returns:

graphviz representation of autograd graph

Return type:

AutoGradDot

analogvnn.utils.render_autograd_graph.make_autograd_obj_from_module(module: torch.nn.Module, *args: torch.Tensor, additional_params: Optional[dict] = None, show_attrs: bool = True, show_saved: bool = True, max_attr_chars: int = 50, from_forward: bool = False, **kwargs: torch.Tensor) AutoGradDot[source]#

Compile Graphviz representation of PyTorch autograd graph from forward pass.

Parameters:
  • module (nn.Module) – PyTorch model

  • *args (Tensor) – input to the model

  • additional_params (dict) – dict of additional params to label nodes with

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

  • from_forward (bool) – if True then use autograd graph otherwise analogvvn graph

  • **kwargs (Tensor) – input to the model

Returns:

graphviz representation of autograd graph

Return type:

AutoGradDot

analogvnn.utils.render_autograd_graph.get_autograd_dot_from_trace(trace) graphviz.Digraph[source]#

Produces graphs of torch.jit.trace outputs.

Example: >>> trace, = torch.jit.trace(model, args=(x,)) >>> dot = get_autograd_dot_from_trace(trace)

Parameters:

trace (torch.jit.trace) – the trace object to visualize.

Returns:

the resulting graph.

Return type:

graphviz.Digraph

analogvnn.utils.render_autograd_graph.get_autograd_dot_from_outputs(outputs: Union[torch.Tensor, Sequence[torch.Tensor]], named_params: Union[Dict[str, Any], Iterator[Tuple[str, torch.nn.Parameter]]], additional_params: Optional[dict] = None, show_attrs: bool = True, show_saved: bool = True, max_attr_chars: int = 50) graphviz.Digraph[source]#

Runs and make Graphviz representation of PyTorch autograd graph from output tensors.

Parameters:
  • outputs (Union[Tensor, Sequence[Tensor]]) – output tensor(s) of forward pass

  • named_params (Union[Dict[str, Any], Iterator[Tuple[str, Parameter]]]) – dict of params to label nodes with

  • additional_params (dict) – dict of additional params to label nodes with

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

Returns:

graphviz representation of autograd graph

Return type:

Digraph

analogvnn.utils.render_autograd_graph.get_autograd_dot_from_module(module: torch.nn.Module, *args: torch.Tensor, additional_params: Optional[dict] = None, show_attrs: bool = True, show_saved: bool = True, max_attr_chars: int = 50, from_forward: bool = False, **kwargs: torch.Tensor) graphviz.Digraph[source]#

Runs and make Graphviz representation of PyTorch autograd graph from forward pass.

Parameters:
  • module (nn.Module) – PyTorch model

  • *args (Tensor) – input to the model

  • additional_params (dict) – dict of additional params to label nodes with

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

  • from_forward (bool) – if True then use autograd graph otherwise analogvvn graph

  • **kwargs (Tensor) – input to the model

Returns:

graphviz representation of autograd graph

Return type:

Digraph

analogvnn.utils.render_autograd_graph.save_autograd_graph_from_outputs(filename: Union[str, pathlib.Path], outputs: Union[torch.Tensor, Sequence[torch.Tensor]], named_params: Union[Dict[str, Any], Iterator[Tuple[str, torch.nn.Parameter]]], additional_params: Optional[dict] = None, show_attrs: bool = True, show_saved: bool = True, max_attr_chars: int = 50) str[source]#

Save Graphviz representation of PyTorch autograd graph from output tensors.

Parameters:
  • filename (Union[str, Path]) – filename to save the graph to

  • outputs (Union[Tensor, Sequence[Tensor]]) – output tensor(s) of forward pass

  • named_params (Union[Dict[str, Any], Iterator[Tuple[str, Parameter]]]) – dict of params to label nodes with

  • additional_params (dict) – dict of additional params to label nodes with

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

Returns:

The (possibly relative) path of the rendered file.

Return type:

str

analogvnn.utils.render_autograd_graph.save_autograd_graph_from_module(filename: Union[str, pathlib.Path], module: torch.nn.Module, *args: torch.Tensor, additional_params: Optional[dict] = None, show_attrs: bool = True, show_saved: bool = True, max_attr_chars: int = 50, from_forward: bool = False, **kwargs: torch.Tensor) str[source]#

Save Graphviz representation of PyTorch autograd graph from forward pass.

Parameters:
  • filename (Union[str, Path]) – filename to save the graph to

  • module (nn.Module) – PyTorch model

  • *args (Tensor) – input to the model

  • additional_params (dict) – dict of additional params to label nodes with

  • show_attrs (bool) – whether to display non-tensor attributes of backward nodes (Requires PyTorch version >= 1.9)

  • show_saved (bool) – whether to display saved tensor nodes that are not by custom autograd functions. Saved tensor nodes for custom functions, if present, are always displayed. (Requires PyTorch version >= 1.9)

  • max_attr_chars (int) – if show_attrs is True, sets max number of characters to display for any given attribute.

  • from_forward (bool) – if True then use autograd graph otherwise analogvvn graph

  • **kwargs (Tensor) – input to the model

Returns:

The (possibly relative) path of the rendered file.

Return type:

str

analogvnn.utils.render_autograd_graph.save_autograd_graph_from_trace(filename: Union[str, pathlib.Path], trace) str[source]#

Save Graphviz representation of PyTorch autograd graph from trace.

Parameters:
  • filename (Union[str, Path]) – filename to save the graph to

  • trace (torch.jit.trace) – the trace object to visualize.

Returns:

The (possibly relative) path of the rendered file.

Return type:

str