Function visualize_with_clusters

devtools/visualization/visualization_utils.py:178–279

Visualize exported programs using the Model Explorer. The QDQ clusters and individual partitions are highlighted. To install the Model Explorer, run `devtools/install_requirements.sh`. To display a stored json file, first launch the Model Explorer server by running `model-explorer`,

(
    exported_program: ExportedProgram,
    json_file_name: str | None = None,
    no_open_in_browser: bool = False,
    reuse_server: bool = False,
    get_node_partition_name: Callable[[Node], str | None] = lambda node: node.meta.get(
        "delegation_tag", None
    ),
    get_node_qdq_cluster_name: Callable[
        [Node], str | None
    ] = lambda node: node.meta.get("cluster", None),
    **kwargs,
)

def visualize_with_clusters(
    exported_program: ExportedProgram,
    json_file_name: str | None = None,
    no_open_in_browser: bool = False,
    reuse_server: bool = False,
    get_node_partition_name: Callable[[Node], str | None] = lambda node: node.meta.get(
        "delegation_tag", None
    ),
    get_node_qdq_cluster_name: Callable[
        [Node], str | None
    ] = lambda node: node.meta.get("cluster", None),
    **kwargs,
):
    """Visualize exported programs using the Model Explorer. The QDQ clusters and individual partitions are highlighted.

        To install the Model Explorer, run `devtools/install_requirements.sh`.
        To display a stored json file, first launch the Model Explorer server by running `model-explorer`, and then
         use the GUI to open the json.

        NOTE: FireFox seems to have issues rendering the graphs. Other browsers seem to work well.

    :param exported_program: Program to visualize.
    :param json_file_name: If not None, a JSON of the visualization will be stored in the provided file. The JSON can
                            then be opened in the Model Explorer GUI later.
                           If None, a Model Explorer instance will be launched with the model visualization.
    :param no_open_in_browser: If `True`, a browser instance with the model explorer will NOT be launched, and only the
                                URI to the model explorer server with the visualization will be printed.
    :param reuse_server: If True, an existing instance of the Model Explorer server will be used (if one exists).
                          Otherwise, a new instance at a separate port will start.
    :param get_node_partition_name: Function which takes a `Node` and returns a string with the name of the partition
                                     the `Node` belongs to, or `None` if it has no partition.
    :param get_node_qdq_cluster_name: Function which takes a `Node` and returns a string with the name of the QDQ
                                       cluster the `Node` belongs to, or `None` if it has no cluster.
    :param kwargs: Additional kwargs for the `visualize_from_config()` function.
    """

    cur_config = config()

    # Create a Model Explorer graph from the `exported_program`.
    adapter = PytorchExportedProgramAdapterImpl(
        exported_program, consts.DEFAULT_SETTINGS
    )
    graphs = adapter.convert()

    nodes = list(exported_program.graph.nodes)
    explorer_nodes = graphs["graphs"][0].nodes

    # Highlight QDQ clusters and individual partitions.
    known_partition_names = []
    for explorer_node, node in zip(explorer_nodes, nodes, strict=True):
        # Generate the `namespace` for the node, which will determine node grouping in the visualizer.
        # The character "/" is used as a divider when a node has multiple namespaces.
        namespace = ""

        if (partition_name := get_node_partition_name(node)) is not None:
            # If the nodes are tagged by the partitioner, highlight the tagged groups.

            # Create a custom naming for the partitions ("partition <i>" where `i` = 0, 1, 2, ...).