onnxslim 0.1.80__py3-none-any.whl

onnxslim/utils.py

@@ -0,0 +1,794 @@
+from __future__ import annotations
+
+import importlib.util
+import logging
+import os
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import onnx
+from onnx import checker, helper
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.misc.tabulate import SEPARATING_LINE, tabulate
+from onnxslim.third_party.onnx_graphsurgeon.logger.logger import G_LOGGER
+
+logger = logging.getLogger("onnxslim")
+
+
+import ml_dtypes
+
+try:
+    from onnx._mapping import TensorDtypeMap
+except ImportError:
+    from onnx.mapping import TensorDtypeMap
+
+TENSOR_TYPE_MAP = {}
+
+candidates = [
+    ("BFLOAT16", "bfloat16", "UINT16"),
+    ("FLOAT8E4M3FN", "float8_e4m3fn", "UINT8"),
+    ("FLOAT8E4M3FNUZ", "float8_e4m3fnuz", "UINT8"),
+    ("FLOAT8E5M2", "float8_e5m2", "UINT8"),
+    ("FLOAT8E5M2FNUZ", "float8_e5m2fnuz", "UINT8"),
+    ("UINT4", "uint4", "INT32"),
+    ("INT4", "int4", "INT32"),
+    ("FLOAT4E2M1", "float4_e2m1fn", "UINT8"),
+]
+
+for onnx_name, ml_name, storage_name in candidates:
+    if hasattr(onnx.TensorProto, onnx_name) and hasattr(ml_dtypes, ml_name):
+        TENSOR_TYPE_MAP[int(getattr(onnx.TensorProto, onnx_name))] = TensorDtypeMap(
+            np.dtype(getattr(ml_dtypes, ml_name)),
+            int(getattr(onnx.TensorProto, storage_name)),
+            f"TensorProto.{onnx_name}",
+        )
+
+
+def init_logging(verbose=False):
+    """Configure the logging settings for the application based on the verbosity level."""
+    logger = logging.getLogger("onnxslim")
+
+    if verbose:  # DEBUG
+        logger.setLevel(logging.DEBUG)
+        G_LOGGER.severity = logging.DEBUG
+    else:  # ERROR
+        logger.setLevel(logging.ERROR)
+        G_LOGGER.severity = logging.ERROR
+
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stderr)
+        formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    G_LOGGER.colors = False
+
+    if is_onnxruntime_available():
+        import onnxruntime as ort
+
+        ort.set_default_logger_severity(3)
+
+    return logger
+
+
+def format_bytes(size: int | tuple[int, ...]) -> str:
+    """Convert byte sizes into human-readable format with appropriate units (B, KB, MB, GB)."""
+    if isinstance(size, int):
+        size = (size,)
+    elif isinstance(size, np.integer):
+        size = (int(size),)
+
+    units = ["B", "KB", "MB", "GB"]
+    formatted_sizes = []
+
+    for size_in_bytes in size:
+        unit_index = 0
+        while size_in_bytes >= 1024 and unit_index < len(units) - 1:
+            size_in_bytes /= 1024
+            unit_index += 1
+
+        formatted_size = f"{size_in_bytes:.2f} {units[unit_index]}"
+        formatted_sizes.append(formatted_size)
+
+    if len(formatted_sizes) == 1:
+        return formatted_sizes[0]
+    else:
+        return f"{formatted_sizes[0]} ({formatted_sizes[1]})"
+
+
+def onnx_dtype_to_numpy(onnx_dtype: int) -> np.dtype:
+    """Maps an ONNX dtype to its corresponding NumPy dtype."""
+    tensor_dtype = TENSOR_TYPE_MAP.get(onnx_dtype)
+
+    if tensor_dtype:
+        return tensor_dtype.np_dtype
+
+    if onnx_dtype in onnx.helper.get_all_tensor_dtypes():
+        return np.dtype(helper.tensor_dtype_to_np_dtype(onnx_dtype))
+
+    return "UNDEFINED"
+
+
+def gen_onnxruntime_input_data(
+    model: onnx.ModelProto, model_check_inputs: list[str] | None = None
+) -> dict[str, np.ndarray]:
+    """Generate random input data for an ONNX model considering potential specific input shapes and types."""
+    input_info = {}
+    for input_tensor in model.graph.input:
+        name = input_tensor.name
+        shape = []
+        for dim in input_tensor.type.tensor_type.shape.dim:
+            if dim.HasField("dim_param"):
+                shape.append(dim.dim_param)
+            elif dim.HasField("dim_value"):
+                shape.append(dim.dim_value)
+            else:
+                shape.append(None)
+        dtype = onnx_dtype_to_numpy(input_tensor.type.tensor_type.elem_type)
+
+        input_info[name] = {"shape": shape, "dtype": dtype}
+
+    if model_check_inputs:
+        for model_check_input in model_check_inputs:
+            key, value = model_check_input.rsplit(":", 1)
+            if value.endswith(".npy"):
+                if key not in input_info:
+                    raise Exception(
+                        f"model_check_input name:{key} not found in model, available keys: {' '.join(input_info.keys())}"
+                    )
+                data = np.load(value)
+                input_info[key] = {"data": data}
+            else:
+                values_list = [int(val) for val in value.split(",")]
+                if key in input_info:
+                    input_info[key]["shape"] = values_list
+                else:
+                    raise Exception(
+                        f"model_check_input name:{key} not found in model, available keys: {' '.join(input_info.keys())}"
+                    )
+
+    input_data_dict = {}
+    for name, info in input_info.items():
+        if "data" in info:
+            input_data_dict[name] = info["data"]
+        else:
+            shapes = [shape if (shape != -1 and not isinstance(shape, str)) else 1 for shape in info["shape"]]
+            shapes = shapes or [1]
+            dtype = info["dtype"]
+
+            if dtype in {np.int32, np.int64}:
+                random_data = np.random.randint(10, size=shapes).astype(dtype)
+            else:
+                random_data = np.random.rand(*shapes).astype(dtype)
+            input_data_dict[name] = random_data
+
+    return input_data_dict
+
+
+def onnxruntime_inference(model: onnx.ModelProto, input_data: dict) -> dict[str, np.array]:
+    """Perform inference using ONNX Runtime on the given model and input data."""
+    import os
+    import tempfile
+
+    import onnx
+    import onnxruntime as rt
+
+    if model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF:
+        tmp_dir = tempfile.TemporaryDirectory()
+        tmp_path = os.path.join(tmp_dir.name, "tmp.onnx")
+        location = f"{os.path.basename(tmp_path)}.data"
+        if os.path.exists(location):
+            os.remove(location)
+        onnx.save(
+            model,
+            tmp_path,
+            save_as_external_data=True,
+            all_tensors_to_one_file=True,
+            location=location,
+        )
+        onnx_model = tmp_path
+    else:
+        onnx_model = model.SerializeToString()
+
+    sess = rt.InferenceSession(onnx_model, providers=["CPUExecutionProvider"])
+    onnx_output = sess.run(None, input_data)
+
+    output_names = [output.name for output in sess.get_outputs()]
+    onnx_output = dict(zip(output_names, onnx_output))
+
+    if isinstance(onnx_model, str):
+        model = onnx.load(onnx_model)
+
+    return onnx_output, model
+
+
+def format_model_info(model_info_list: dict | list[dict], elapsed_time: float | None = None):
+    assert model_info_list, "model_info_list must contain more than one model info"
+    from colorama import Fore, init
+
+    init()
+    if not isinstance(model_info_list, (list, tuple)):
+        model_info_list = [model_info_list]
+
+    final_op_info = []
+    final_op_info.extend(
+        (
+            ["Model Name"] + [item.tag for item in model_info_list],
+            [SEPARATING_LINE] * (len(model_info_list) + 1),
+            ["Model Info"]
+            + ["Op Set: " + item.op_set + " / IR Version: " + item.ir_version for item in model_info_list],
+            [SEPARATING_LINE] * (len(model_info_list) + 1),
+        )
+    )
+
+    def get_io_info(model_info_list, tag=None):
+        if tag == "OUT":
+            ios = [op_type for model_info in model_info_list for op_type in model_info.output_info]
+        else:
+            ios = [op_type for model_info in model_info_list for op_type in model_info.input_info]
+        ios = list(dict.fromkeys([io.name for io in ios]))
+        io_info = []
+        for io in ios:
+            input_info_list = [f"{tag}: {io}"]
+            for model_info in model_info_list:
+                if tag == "OUT":
+                    io_tensor = model_info.output_maps.get(io, None)
+                else:
+                    io_tensor = model_info.input_maps.get(io, None)
+                inputs_shape = (io_tensor.dtype, io_tensor.shape) if io_tensor else ""
+                if isinstance(inputs_shape, (list, tuple)):
+                    inputs_shape = ": ".join([str(i) for i in inputs_shape])
+                input_info_list.append(inputs_shape)
+            io_info.append(input_info_list)
+
+        return io_info
+
+    final_op_info.extend(get_io_info(model_info_list, "IN"))
+    final_op_info.extend(get_io_info(model_info_list, "OUT"))
+
+    final_op_info.append([SEPARATING_LINE] * (len(model_info_list) + 1))
+
+    all_ops = {op_type for model_info in model_info_list for op_type in model_info.op_type_counts}
+    sorted_ops = sorted(all_ops)
+    for op in sorted_ops:
+        op_info_list = [op]
+        float_number = model_info_list[0].op_type_counts.get(op, 0)
+        op_info_list.append(float_number)
+        for model_info in model_info_list[1:]:
+            slimmed_number = model_info.op_type_counts.get(op, 0)
+            if float_number > slimmed_number:
+                slimmed_number = Fore.GREEN + str(slimmed_number) + Fore.WHITE
+            op_info_list.append(slimmed_number)
+
+        final_op_info.append(op_info_list)
+    final_op_info.extend(
+        (
+            [SEPARATING_LINE] * (len(model_info_list) + 1),
+            ["Model Size"] + [format_bytes(model_info.model_size) for model_info in model_info_list],
+        )
+    )
+    if elapsed_time:
+        final_op_info.extend(
+            (
+                [SEPARATING_LINE] * (len(model_info_list) + 1),
+                ["Elapsed Time", f"{elapsed_time:.2f} s"],
+            )
+        )
+
+    return final_op_info
+
+
+def print_model_info_as_table(model_info_list: dict | list[dict], elapsed_time: float | None = None):
+    """Prints the model information as a formatted table for the given model name and list of model details."""
+    if not isinstance(model_info_list, (list, tuple)):
+        model_info_list = [model_info_list]
+
+    final_op_info = format_model_info(model_info_list, elapsed_time)
+    lines = tabulate(
+        final_op_info,
+        headers=[],
+        tablefmt="pretty",
+        maxcolwidths=[None] + [40] * len(model_info_list),
+    ).split("\n")
+    if elapsed_time:
+        time_row = lines[-2].split("|")
+        time_row[-3] = (
+            time_row[-2][: len(time_row[-2]) // 2 + 1] + time_row[-3] + time_row[-2][len(time_row[-2]) // 2 :]
+        )
+        time_row.pop(-2)
+        lines[-2] = "|".join(time_row)
+    output = "\n".join([line if line != "| \x01 |" else lines[0] for line in lines])
+
+    print(output)
+
+
+def dump_model_info_to_disk(model_info: dict):
+    """Writes model information to a CSV file for a given model name and dictionary of model info."""
+    import csv
+
+    csv_file_path = f"{model_info.tag}_model_info.csv"
+    with open(csv_file_path, "a", newline="") as csvfile:  # Use 'a' for append mode
+        fieldnames = ["NodeName", "OpType", "OutputDtype", "OutputShape"]
+        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+
+        # If the file is empty, write the header
+        if csvfile.tell() == 0:
+            writer.writeheader()
+
+        # Write the data
+        for node_name, info in model_info.op_info.items():
+            op_type, output_info_list = info.op, info.outputs
+            if len(output_info_list) >= 1:
+                # Write the first row with actual NodeName and OpType
+                row_data_first = {
+                    "NodeName": node_name,
+                    "OpType": op_type,
+                    "OutputDtype": output_info_list[0].dtype,  # First entry in the list
+                    "OutputShape": output_info_list[0].shape,  # First entry in the list
+                }
+                writer.writerow(row_data_first)
+
+                # Write subsequent rows with empty strings for NodeName and OpType
+                for output_dtype, output_shape in output_info_list[1:]:
+                    row_data_empty = {
+                        "NodeName": "",
+                        "OpType": "",
+                        "OutputDtype": output_dtype,
+                        "OutputShape": output_shape,
+                    }
+                    writer.writerow(row_data_empty)
+    print(f"Model info written to {csv_file_path}")
+
+
+def get_opset(model: onnx.ModelProto) -> int:
+    """Returns the ONNX opset version for a given model."""
+    try:
+        for importer in model.opset_import:
+            if importer.domain in {"", "ai.onnx"}:
+                return importer.version
+
+        return None
+    except Exception:
+        return None
+
+
+def get_ir_version(model: onnx.ModelProto) -> int:
+    """Returns the ONNX ir version for a given model."""
+    try:
+        return model.ir_version
+    except Exception:
+        return None
+
+
+class TensorInfo:
+    def __init__(self, tensor):
+        self.dtype: np.dtype = np.float32
+        self.shape: tuple[str | int] = None
+
+        self._extract_info(tensor)
+
+    def _extract_info(self, tensor):
+        """Extract the data type and shape of an ONNX tensor."""
+        self.dtype = onnx_dtype_to_numpy(tensor.type.tensor_type.elem_type)
+        shape = None
+        if tensor.type.tensor_type.HasField("shape"):
+            shape = []
+            for dim in tensor.type.tensor_type.shape.dim:
+                if dim.HasField("dim_param"):
+                    shape.append(dim.dim_param)
+                elif dim.HasField("dim_value"):
+                    shape.append(dim.dim_value)
+                else:
+                    shape.append("?")
+
+        self.shape = tuple(shape) if shape is not None else None
+        self.name = tensor.name
+
+
+class OperatorInfo:
+    def __init__(self, operator, outputs=None):
+        self.name: str = None
+        self.op: str = None
+
+        self._extract_info(operator)
+        self.outputs = outputs
+
+    def _extract_info(self, operator):
+        self.name: str = operator.name
+        self.op: str = operator.op_type
+
+
+class ModelInfo:
+    def __init__(self, model: str | onnx.ModelProto, tag: str = "OnnxSlim"):
+        if isinstance(model, str):
+            tag = Path(model).name
+            model = onnx.load(model)
+
+        self.tag: str = tag
+        self.model_size: int = -1
+        self.op_set: str = None
+        self.ir_version: str = None
+        self.op_type_counts: dict[str, int] = defaultdict(int)
+        self.op_info: dict[str, dict] = {}
+        self.input_info: list[str, tuple[str, tuple]] = []
+        self.output_info: list[str, tuple[str, tuple]] = []
+
+        self._summarize_model(model)
+
+    def _summarize_model(self, model):
+        self.op_set = str(get_opset(model))
+        self.ir_version = str(get_ir_version(model))
+        self.model_size = get_initializer_size(model)
+
+        for input in model.graph.input:
+            self.input_info.append(TensorInfo(input))
+
+        for output in model.graph.output:
+            self.output_info.append(TensorInfo(output))
+
+        value_info_dict = {value_info.name: value_info for value_info in model.graph.value_info}
+
+        def get_graph_node_info(graph: onnx.GraphProto) -> dict[str, list[str]]:
+            for node in graph.node:
+                op_type = node.op_type
+                self.op_type_counts[op_type] += 1
+                output_tensor_info = []
+                for output in node.output:
+                    if output in value_info_dict:
+                        tensor = value_info_dict[output]
+                        tensor_info = TensorInfo(tensor)
+                        output_tensor_info.append(tensor_info)
+
+                self.op_info[node.name] = OperatorInfo(node, output_tensor_info)
+
+                for attr in node.attribute:
+                    ATTR_TYPE_MAPPING = {v: k for k, v in onnx.AttributeProto.AttributeType.items()}
+                    if attr.type in ATTR_TYPE_MAPPING:
+                        attr_str = ATTR_TYPE_MAPPING[attr.type]
+                        if attr_str == "GRAPH":
+                            get_graph_node_info(attr.g)
+
+        get_graph_node_info(model.graph)
+
+    @property
+    def input_maps(self):
+        self.input_dict = {input_info.name: input_info for input_info in self.input_info}
+
+        return self.input_dict
+
+    @property
+    def output_maps(self):
+        self.output_dict = {output_info.name: output_info for output_info in self.output_info}
+
+        return self.output_dict
+
+
+def summarize_model(model: str | onnx.ModelProto, tag="OnnxModel") -> dict:
+    """Generates a summary of the ONNX model, including model size, operations, and tensor shapes."""
+    logger.debug("Start summarizing model.")
+    model_info = ModelInfo(model, tag)
+    logger.debug("Finish summarizing model.")
+    return model_info
+
+
+def model_save_as_external_data(model: onnx.ModelProto, model_path: str):
+    """Save an ONNX model with tensor data as an external file."""
+    location = f"{os.path.basename(model_path)}.data"
+    if os.path.exists(location):
+        os.remove(location)
+    onnx.save(
+        model,
+        model_path,
+        save_as_external_data=True,
+        all_tensors_to_one_file=True,
+        location=location,
+    )
+
+
+def check_onnx(model: onnx.ModelProto, model_check_inputs=None):
+    """Validates an ONNX model by generating input data and performing inference to check outputs."""
+    input_data_dict = gen_onnxruntime_input_data(model, model_check_inputs)
+    raw_onnx_output, model = onnxruntime_inference(model, input_data_dict)
+
+    return input_data_dict, raw_onnx_output, model
+
+
+def check_point(model: onnx.ModelProto):
+    """Imports an ONNX model checkpoint into a Graphsurgeon graph representation."""
+    return gs.import_onnx(model)
+
+
+def save(
+    model: onnx.ModelProto,
+    model_path: str,
+    model_check: bool = False,
+    save_as_external_data: bool = False,
+    model_info: dict | None = None,
+):
+    """Save an ONNX model to a specified path, with optional model checking for validity."""
+    if model_check:
+        try:
+            checker.check_model(model)
+        except ValueError:
+            logger.warning("Model too large and cannot be checked.")
+
+    if model_path:  # model larger than 2GB can be saved, but compiler like trtexec won't parse it
+        if get_initializer_size(model) <= checker.MAXIMUM_PROTOBUF and not save_as_external_data:
+            onnx.save(model, model_path)
+        else:
+            import os
+
+            location = f"{os.path.basename(model_path)}.data"
+            if os.path.exists(location):
+                os.remove(location)
+            onnx.save(
+                model,
+                model_path,
+                save_as_external_data=True,
+                all_tensors_to_one_file=True,
+                location=location,
+            )
+            logger.debug("Model too large and saved as external data automatically.")
+
+            if model_info:
+                model_size = model.ByteSize()
+                model_info.model_size = [model_size, model_info.model_size]
+
+
+def check_result(raw_onnx_output, slimmed_onnx_output):
+    """Verify the consistency of outputs between the raw and slimmed ONNX models, logging warnings if discrepancies are
+    detected.
+    """
+    if set(raw_onnx_output.keys()) != set(slimmed_onnx_output.keys()):
+        print("Model output mismatch after slimming.")
+        print(f"Raw model output keys: {raw_onnx_output.keys()}")
+        print(f"Slimmed model output keys: {slimmed_onnx_output.keys()}")
+        print("Please check the model carefully.")
+        return False
+    else:
+        for key in raw_onnx_output.keys():
+            if not np.allclose(
+                raw_onnx_output[key],
+                slimmed_onnx_output[key],
+                rtol=1e-03,
+                atol=1e-04,
+                equal_nan=True,
+            ):
+                print(f"\033[31mModel output {key} mismatch after slimming.")
+                print("\033[31mPlease check the model carefully.")
+                return False
+
+    return True
+
+
+def get_numpy_type(onnx_type):
+    if not isinstance(onnx_type, int):
+        # Already a NumPy type
+        return onnx_type
+
+    numpy_unsupported_types = [
+        onnx.TensorProto.BFLOAT16,
+        onnx.TensorProto.FLOAT8E4M3FN,
+        onnx.TensorProto.FLOAT8E4M3FNUZ,
+        onnx.TensorProto.FLOAT8E5M2,
+        onnx.TensorProto.FLOAT8E5M2FNUZ,
+    ]
+
+    # TENSOR_TYPE_TO_NP_TYPE maps types unsupported by NumPy to random other types.
+    # This obviously breaks things, so we need to treat this as a special case.
+    if onnx_type not in numpy_unsupported_types and onnx_type in onnx.helper.get_all_tensor_dtypes():
+        return onnx.helper.tensor_dtype_to_np_dtype(onnx_type)
+    return None
+
+
+def get_itemsize(dtype):
+    np_dtype = get_numpy_type(dtype)
+    if np_dtype is not None:
+        return np.dtype(np_dtype).itemsize
+
+    if dtype == onnx.TensorProto.BFLOAT16:
+        return 2
+
+    if dtype in [
+        onnx.TensorProto.FLOAT8E4M3FN,
+        onnx.TensorProto.FLOAT8E4M3FNUZ,
+        onnx.TensorProto.FLOAT8E5M2,
+        onnx.TensorProto.FLOAT8E5M2FNUZ,
+    ]:
+        return 1
+
+    print(f"Unknown ONNX dtype: {dtype}")
+    raise ValueError(f"Unsupported TensorProto dtype: {dtype}")
+
+
+def calculate_tensor_size(tensor):
+    """Calculates the size of an ONNX tensor in bytes based on its shape and data type size."""
+    shape = tensor.dims
+    num_elements = np.prod(shape) if shape else 0
+    element_size = get_itemsize(tensor.data_type)
+    return num_elements * element_size
+
+
+def get_initializer_size(model):
+    """Calculate total size of all subgraphs in an ONNX model."""
+    total_size = get_graph_initializer_size(model.graph)
+    return total_size
+
+
+def get_graph_initializer_size(graph):
+    initializer_size = 0
+    for tensor in graph.initializer:
+        tensor_size = calculate_tensor_size(tensor)
+        initializer_size += tensor_size
+
+    for node in graph.node:
+        if node.op_type == "Constant":
+            for attr in node.attribute:
+                if attr.name == "value" and attr.type == onnx.AttributeProto.TENSOR:
+                    initializer_size += calculate_tensor_size(attr.t)
+
+        elif node.op_type == "If":
+            initializer_size += get_graph_initializer_size(node.attribute[0].g)
+            initializer_size += get_graph_initializer_size(node.attribute[1].g)
+        elif node.op_type == "Loop":
+            initializer_size += get_graph_initializer_size(node.attribute[0].g)
+        elif node.op_type == "Scan":
+            initializer_size += get_graph_initializer_size(node.attribute[0].g)
+
+    return initializer_size
+
+
+def is_onnxruntime_available():
+    if importlib.util.find_spec("onnxruntime") is None:
+        logger = logging.getLogger("onnxslim")
+        logger.debug("onnxruntime is not available, please install it first for better optimization")
+        return False
+    else:
+        try:
+            # in case of onnxruntime import error
+            import onnxruntime as ort
+
+            if hasattr(ort, "__version__"):
+                return True
+            else:
+                return False
+        except:
+            logger = logging.getLogger("onnxslim")
+            logger.debug("onnxruntime is not available, please install it first for better optimization")
+            return False
+
+
+def check_onnx_compatibility():
+    """Ensure ONNX Runtime and ONNX versions are compatible for model inference."""
+    compatibility_dict = {
+        "1.20": "1.16",
+        "1.19": "1.16",
+        "1.18": "1.16",
+        "1.17": "1.15",
+        "1.16": "1.14.1",
+        "1.15": "1.14",
+        "1.14": "1.13",
+        "1.13": "1.12",
+        "1.12": "1.12",
+        "1.11": "1.11",
+        "1.10": "1.10",
+        "1.9": "1.10",
+        "1.8": "1.9",
+        "1.7": "1.8",
+        "1.6": "1.8",
+        "1.5": "1.7",
+        "1.4": "1.7",
+        "1.3": "1.7",
+        "1.2": "1.6",
+        "1.1": "1.6",
+        "1.0": "1.6",
+        "0.5": "1.5",
+        "0.4": "1.5",
+        "0.3": "1.4",
+        "0.2": "1.3",
+        "0.1": "1.3",
+    }
+    import onnx
+    import onnxruntime
+
+    onnx_version = onnx.__version__
+    # ort_version = onnxruntime.__version__
+    ort_version = ".".join(onnxruntime.__version__.split("+")[0].split(".")[:2])
+    # Check compatibility
+    expected_onnx_version = compatibility_dict.get(ort_version)
+    if expected_onnx_version is None:
+        print(
+            f"Warning: Onnx Runtime version {ort_version} has no specified compatible ONNX version. Compatibility issues may occur."
+        )
+    elif expected_onnx_version == ".".join(onnx_version.split("+")[0].split(".")[:2]):
+        logger.info(
+            f"Installed Onnx Runtime version {ort_version} is compatible with installed ONNX version {onnx_version}."
+        )
+    else:
+        print(
+            f"Warning: Installed Onnx Runtime version {ort_version} is not compatible with installed ONNX version {onnx_version}. Expected ONNX version: {expected_onnx_version}."
+        )
+
+
+def get_max_tensor(model, topk=5):
+    graph = gs.import_onnx(model)
+
+    tensor_map = graph.tensors()
+    constant_tensors = [tensor for tensor in tensor_map.values() if isinstance(tensor, gs.Constant)]
+
+    sub_graphs = graph.subgraphs(recursive=True)
+    sub_graphs_constant_tensors = [
+        [tensor for name, tensor in sub_graph.tensors().items() if isinstance(tensor, gs.Constant)]
+        for sub_graph in sub_graphs
+    ]
+
+    constant_tensors.extend([tensor for tensors in sub_graphs_constant_tensors for tensor in tensors])
+
+    sizes = [tensor.values.size for tensor in constant_tensors]
+    sorted_indices = np.argsort(sizes)[::-1][:topk]
+
+    for i in sorted_indices:
+        tensor = constant_tensors[i]
+        print(
+            f"Tensor name: {tensor.name}, shape: {tensor.values.shape}, dtype: {tensor.values.dtype} size: {tensor.values.size}"
+        )
+
+
+# copied from https://onnx.ai/onnx/api/tools.html
+def update_outputs_dims(
+    model,
+    output_dims,
+):
+    dim_param_set: set[str] = set()
+
+    def init_dim_param_set(dim_param_set, value_infos):
+        for info in value_infos:
+            shape = info.type.tensor_type.shape
+            for dim in shape.dim:
+                if dim.HasField("dim_param"):
+                    dim_param_set.add(dim.dim_param)  # type: ignore
+
+    init_dim_param_set(dim_param_set, model.graph.output)  # type: ignore
+
+    def update_dim(tensor, dim, j, name) -> None:
+        dim_proto = tensor.type.tensor_type.shape.dim[j]
+
+        # if it's int in model, it won't be replaced by original symbol
+        if dim_proto.HasField("dim_value"):
+            return
+
+        if isinstance(dim, int):
+            if dim >= 0:
+                if dim_proto.HasField("dim_value") and dim_proto.dim_value != dim:
+                    raise ValueError(
+                        f"Unable to set dimension value to {dim} for axis {j} of {name}. Contradicts existing dimension value {dim_proto.dim_value}."
+                    )
+                dim_proto.dim_value = dim
+            else:
+                generated_dim_param = name + "_" + str(j)
+                if generated_dim_param in dim_param_set:
+                    raise ValueError(
+                        f"Unable to generate unique dim_param for axis {j} of {name}. Please manually provide a dim_param value."
+                    )
+                dim_proto.dim_param = generated_dim_param
+        elif isinstance(dim, str):
+            dim_proto.dim_param = dim
+        else:
+            raise ValueError(f"Only int or str is accepted as dimension value, incorrect type: {type(dim)}")
+
+    for output in model.graph.output:
+        output_name = output.name
+        output_dim_arr = output_dims[output_name]
+        if output_dim_arr is None:
+            continue
+
+        if len(output.type.tensor_type.shape.dim) == 0:
+            for _ in range(len(output_dim_arr)):
+                output.type.tensor_type.shape.dim.add()
+        for j, dim in enumerate(output_dim_arr):
+            update_dim(output, dim, j, output_name)
+
+    return model