Function visualize

imperative/python/megengine/tools/network_visualize.py:32–276

r"""Load megengine dumped model and visualize graph structure with tensorboard log files. Can also record and print model's statistics like :func:`~.module_stats` Args: model_path: dir path for megengine dumped model. log_path: dir path for tensorboard graph log. input: us

(
    model_path: str,
    log_path: str,
    input: np.ndarray = None,
    inp_dict: dict = None,
    cal_params: bool = True,
    cal_flops: bool = True,
    cal_activations: bool = True,
    logging_to_stdout: bool = True,
    bar_length_max: int = 20,
)

def visualize(
    model_path: str,
    log_path: str,
    input: np.ndarray = None,
    inp_dict: dict = None,
    cal_params: bool = True,
    cal_flops: bool = True,
    cal_activations: bool = True,
    logging_to_stdout: bool = True,
    bar_length_max: int = 20,
):
    r"""Load megengine dumped model and visualize graph structure with tensorboard log files.
    Can also record and print model's statistics like :func:`~.module_stats`

    Args:
      model_path: dir path for megengine dumped model.
      log_path: dir path for tensorboard graph log.
      input: user defined input data for running model and calculating stats,
        alternative with inp_dict, used when the model has only one input.
      inp_dict: input dict for running model and calculating stats, alternative with
        input, used when the model has more than one input.
        When both input and inp_dict are None, a random input will be used.
      cal_params: whether calculate and record params size.
      cal_flops: whether calculate and record op flops.
      cal_activations: whether calculate and record op activations.
      logging_to_stdout: whether print all calculated statistic details.
      bar_length_max: size of bar indicating max flops or parameter size in net stats.
      model_path: str:
      log_path: str:
      input: np.ndarray:
      inp_dict: dict:
      cal_params: bool:
      cal_flops: bool:
      cal_activations: bool:
      logging_to_stdout: bool:
      bar_length_max: int:
    """
    if log_path:
        try:
            from tensorboard.compat.proto.attr_value_pb2 import AttrValue
            from tensorboard.compat.proto.config_pb2 import RunMetadata
            from tensorboard.compat.proto.graph_pb2 import GraphDef
            from tensorboard.compat.proto.node_def_pb2 import NodeDef
            from tensorboard.compat.proto.step_stats_pb2 import (
                AllocatorMemoryUsed,
                DeviceStepStats,
                NodeExecStats,
                StepStats,
            )
            from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
            from tensorboard.compat.proto.versions_pb2 import VersionDef
            from tensorboardX import SummaryWriter
        except ImportError:
            logger.error(
                "TensorBoard and TensorboardX are required for visualize.",
                exc_info=True,
            )
            return