Function _save_model_as_json

devtools/visualization/visualization_utils.py:150–175

Save the graphs stored in the `cur_config` in JSON format, which can be loaded by the Model Explorer GUI. :param cur_config: ModelExplorerConfig containing the graph for visualization. :param file_name: Name of the JSON file for storage.

(cur_config: ModelExplorerConfig, file_name: str)

def _save_model_as_json(cur_config: ModelExplorerConfig, file_name: str):
    """Save the graphs stored in the `cur_config` in JSON format, which can be loaded by the Model Explorer GUI.

    :param cur_config: ModelExplorerConfig containing the graph for visualization.
    :param file_name: Name of the JSON file for storage.
    """
    # Extract the graphs from the config file.
    graphs_list = json.loads(cur_config.get_transferrable_data()["graphs_list"])
    graphs = json.loads(graphs_list[0])["graphs"]

    # The returned dictionary is missing the `collectionLabel` entry. Add it manually.
    for graph in graphs:
        graph["collectionLabel"] = "Executorch"

    # Create the JSON according to the structure required by the Model Explorer GUI.
    json_data = {
        "label": "Executorch",
        "graphs": graphs,
        "graphsWithLevel": [
            {"graph": graph, "level": level} for level, graph in enumerate(graphs)
        ],
    }

    # Store the JSON.
    with open(file_name, "w") as f:
        json.dump(json_data, f)


def visualize_with_clusters(