node_data_for_model_explorer.py

def save_node_data_for_model_explorer(verification_infos: Collection[VerificationInfo], node_names: list[str], model_name: str = "model"
):
    # https://github.com/google-ai-edge/model-explorer/wiki/4.-API-Guide#create-custom-node-data
    # This API is unstable and may change in the future.
    from model_explorer import node_data_builder as ndb

    for field in ("max_abs_diff", "max_rel_diff"):
        # Populate values for the main graph in a model.
        main_graph_results: dict[str, ndb.NodeDataResult] = {}
        for info, node_name in zip(verification_infos, node_names):
            main_graph_results[f"[value] {info.name}"] = ndb.NodeDataResult(
                value=getattr(info, field)
            )
            main_graph_results[node_name] = ndb.NodeDataResult(
                value=getattr(info, field)
            )

        thresholds: list[ndb.ThresholdItem] = [
            ndb.ThresholdItem(value=0.00001, bgColor="#388e3c"),
            ndb.ThresholdItem(value=0.0001, bgColor="#8bc34a"),
            ndb.ThresholdItem(value=0.001, bgColor="#c8e6c9"),
            ndb.ThresholdItem(value=0.01, bgColor="#ffa000"),
            ndb.ThresholdItem(value=1, bgColor="#ff5722"),
            ndb.ThresholdItem(value=100, bgColor="#d32f2f"),
        ]

        # Construct the data for the main graph.
        main_graph_data = ndb.GraphNodeData(
            results=main_graph_results, thresholds=thresholds
        )

        # Construct the data for the model.
        # "main_graph" is defined in _core.py
        model_data = ndb.ModelNodeData(graphsData={"main_graph": main_graph_data})

        model_data.save_to_file(f"{model_name}_{field}.json")