ls_mlkit.util.cgraph module

ls_mlkit.util.cgraph.get_compute_graph(model: Module, input_shape=None, input: dict = None, dir: str = 'compute_graph', filename: str = 'simple_net_graph', format: str = 'pdf') → None[source]

Generate the computing graph of model (format default is pdf)

Parameters:

  • model (torch.nn.Module) – the model to generate the computing graph

  • input_shape (tuple, optional) – the shape of the input. Defaults to None.

  • input (dict, optional) – the input of the model. Defaults to None.

  • dir (str, optional) – the directory to save the computing graph. Defaults to “compute_graph”.

  • filename (str, optional) – the name of the file to save the computing graph. Defaults to “simple_net_graph”.

  • format (str, optional) – the format of the computing graph. Defaults to “pdf”.

Returns:

None