lemonade-sdk 7.0.0__py3-none-any.whl

lemonade/common/onnx_helpers.py

@@ -0,0 +1,176 @@
+"""
+Helper functions for dealing with ONNX files and ONNX models
+"""
+
+import os
+from typing import Tuple, Union
+import re
+import math
+import numpy as np
+import onnx
+import onnxruntime as ort
+import lemonade.common.exceptions as exp
+from lemonade.state import State
+import lemonade.common.build as build
+
+
+def check_model(onnx_file, success_message, fail_message) -> bool:
+    if os.path.isfile(onnx_file):
+        print(success_message)
+    else:
+        print(fail_message)
+        return False
+    try:
+        onnx.checker.check_model(onnx_file)
+        print("\tSuccessfully checked onnx file")
+        return True
+    except onnx.checker.ValidationError as e:
+        print("\tError while checking generated ONNX file")
+        print(e)
+        return False
+
+
+def original_inputs_file(cache_dir: str, build_name: str):
+    return os.path.join(build.output_dir(cache_dir, build_name), "inputs.npy")
+
+
+def onnx_dir(state: State):
+    return os.path.join(build.output_dir(state.cache_dir, state.build_name), "onnx")
+
+
+def get_output_names(
+    onnx_model: Union[str, onnx.ModelProto],
+):  # pylint: disable=no-member
+    # Get output names of ONNX file/model
+    if not isinstance(onnx_model, onnx.ModelProto):  # pylint: disable=no-member
+        onnx_model = onnx.load(onnx_model)
+    return [node.name for node in onnx_model.graph.output]  # pylint: disable=no-member
+
+
+def parameter_count(model):
+    weights = model.graph.initializer
+    parameter_count = 0
+
+    for w in weights:
+        weight = onnx.numpy_helper.to_array(w)
+        parameter_count += np.prod(weight.shape)
+    return parameter_count
+
+
+def io_bytes(onnx_path: str) -> Tuple[int, int]:
+    """Return the number of bytes of each of the inputs and outputs"""
+    # pylint: disable = no-member
+
+    def elem_type_to_bytes(elem_type) -> int:
+        """
+        Convert ONNX's elem_type to the number of bytes used by
+        hardware to send that specific datatype through PCIe
+        """
+        if (
+            elem_type == onnx.TensorProto.DataType.UINT8
+            or elem_type == onnx.TensorProto.DataType.INT8
+            or elem_type == onnx.TensorProto.DataType.BOOL
+        ):
+            # Each bool requires an entire byte
+            return 1
+        elif (
+            elem_type == onnx.TensorProto.DataType.UINT16
+            or elem_type == onnx.TensorProto.DataType.INT16
+            or elem_type == onnx.TensorProto.DataType.FLOAT16
+        ):
+            return 2
+        if (
+            elem_type == onnx.TensorProto.DataType.FLOAT
+            or elem_type == onnx.TensorProto.DataType.INT32
+            or elem_type == onnx.TensorProto.DataType.INT64
+            or elem_type == onnx.TensorProto.DataType.DOUBLE
+            or elem_type == onnx.TensorProto.DataType.UINT64
+        ):
+            # 64 bit ints are treated as 32 bits everywhere
+            # Doubles are treated as floats
+            return 4
+        elif (
+            elem_type == onnx.TensorProto.DataType.COMPLEX64
+            or elem_type == onnx.TensorProto.DataType.COMPLEX128
+            or elem_type == onnx.TensorProto.DataType.STRING
+            or elem_type == onnx.TensorProto.DataType.UNDEFINED
+        ):
+            raise exp.Error("Unsupported data type")
+        else:
+            raise exp.Error("Unsupported data type (unknown to ONNX)")
+
+    def get_nodes_bytes(nodes):
+        nodes_bytes = {}
+        for node in nodes:
+
+            # Get the number of the data type
+            dtype_bytes = elem_type_to_bytes(node.type.tensor_type.elem_type)
+
+            # Calculate the total number of elements based on the shape
+            shape = str(node.type.tensor_type.shape.dim)
+            num_elements = np.prod([int(s) for s in shape.split() if s.isdigit()])
+
+            # Assign a total number of bytes to each node
+            nodes_bytes[node.name] = num_elements * dtype_bytes
+
+        return nodes_bytes
+
+    # Get the number of bytes of each of the inputs and outputs
+    model = onnx.load(onnx_path)
+    onnx_input_bytes = get_nodes_bytes(model.graph.input)
+    onnx_output_bytes = get_nodes_bytes(model.graph.output)
+
+    return int(sum(onnx_input_bytes.values())), int(sum(onnx_output_bytes.values()))
+
+
+def dtype_ort2str(dtype_str: str):
+    if dtype_str == "float16":
+        datatype = "float16"
+    elif dtype_str == "float":
+        datatype = "float32"
+    elif dtype_str == "double":
+        datatype = "float64"
+    elif dtype_str == "long":
+        datatype = "int64"
+    else:
+        datatype = dtype_str
+    return datatype
+
+
+def dummy_inputs(onnx_file: str) -> dict:
+    # Generate dummy inputs of the expected shape and type for the input model
+    sess_options = ort.SessionOptions()
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    onnx_session = ort.InferenceSession(onnx_file, sess_options)
+    sess_input = onnx_session.get_inputs()
+
+    input_stats = []
+    for _idx, input_ in enumerate(range(len(sess_input))):
+        input_name = sess_input[input_].name
+        input_shape = sess_input[input_].shape
+
+        # TODO: Use onnx update_inputs_outputs_dims to automatically freeze models
+        for dim in input_shape:
+            if isinstance(dim, str) is True or math.isnan(dim) is True:
+                raise AssertionError(
+                    "Error: Model has dynamic inputs. Freeze the graph and try again"
+                )
+
+        input_type = sess_input[input_].type
+        input_stats.append([input_name, input_shape, input_type])
+
+    input_feed = {}
+    for stat in input_stats:
+        dtype_str = re.search(r"\((.*)\)", stat[2])
+        assert dtype_str is not None
+        datatype = dtype_ort2str(dtype_str.group(1))
+        input_feed[stat[0]] = np.random.rand(*stat[1]).astype(datatype)
+    return input_feed
+
+
+def get_opset(model: onnx.ModelProto) -> int:
+    return getattr(model.opset_import[0], "version", None)
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/common/plugins.py

@@ -0,0 +1,10 @@
+import pkgutil
+import importlib
+
+
+def discover():
+    return {
+        name: importlib.import_module(name)
+        for _, name, _ in pkgutil.iter_modules()
+        if name.startswith("turnkeyml_plugin_")
+    }

lemonade/common/printing.py

@@ -0,0 +1,110 @@
+import os
+import re
+import enum
+import sys
+import math
+
+
+class Colors:
+    HEADER = "\033[95m"
+    OKBLUE = "\033[94m"
+    OKCYAN = "\033[96m"
+    OKGREEN = "\033[92m"
+    WARNING = "\033[93m"
+    FAIL = "\033[91m"
+    ENDC = "\033[0m"
+    BOLD = "\033[1m"
+    UNDERLINE = "\033[4m"
+
+
+def log(txt, c=Colors.ENDC, end="", is_error=False):
+    logn(txt, c=c, end=end, is_error=is_error)
+
+
+def logn(txt, c=Colors.ENDC, end="\n", is_error=False):
+    file = sys.stderr if is_error else sys.stdout
+    print(c + txt + Colors.ENDC, end=end, flush=True, file=file)
+
+
+class LogType(enum.Enum):
+    ERROR = "Error:"
+    SUCCESS = "Woohoo!"
+    WARNING = "Warning:"
+    INFO = "Info:"
+
+
+def clean_print(type: LogType, msg):
+    # Replace path to user’s home directory by a tilde symbol (~)
+    home_directory = os.path.expanduser("~")
+    home_directory_escaped = re.escape(home_directory)
+    msg = re.sub(home_directory_escaped, "~", msg)
+
+    # Split message into list, remove leading spaces and line breaks
+    msg = msg.split("\n")
+    msg = [line.lstrip() for line in msg]
+    while msg[0] == "" and len(msg) > 1:
+        msg.pop(0)
+
+    # Print message
+    indentation = len(type.value) + 1
+    if type == LogType.ERROR:
+        log(f"\n{type.value} ".rjust(indentation), c=Colors.FAIL, is_error=True)
+    elif type == LogType.SUCCESS:
+        log(f"\n{type.value} ".rjust(indentation), c=Colors.OKGREEN)
+    elif type == LogType.WARNING:
+        log(f"\n{type.value} ".rjust(indentation), c=Colors.WARNING)
+    elif type == LogType.INFO:
+        log(f"\n{type.value} ".rjust(indentation), c=Colors.OKCYAN)
+
+    is_error = type == LogType.ERROR
+    for line_idx, line in enumerate(msg):
+        if line_idx != 0:
+            log(" " * indentation)
+        s_line = line.split("**")
+        for idx, l in enumerate(s_line):
+            c = Colors.ENDC if idx % 2 == 0 else Colors.BOLD
+            if idx != len(s_line) - 1:
+                log(l, c=c, is_error=is_error)
+            else:
+                logn(l, c=c, is_error=is_error)
+
+
+def log_error(msg):
+    clean_print(LogType.ERROR, str(msg))
+    # ASCII art credit:
+    # https://textart4u.blogspot.com/2014/05/the-fail-whale-ascii-art-code.html
+    logn(
+        """\n▄██████████████▄▐█▄▄▄▄█▌
+██████▌▄▌▄▐▐▌███▌▀▀██▀▀
+████▄█▌▄▌▄▐▐▌▀███▄▄█▌
+▄▄▄▄▄██████████████\n\n""",
+        is_error=True,
+    )
+
+
+def log_success(msg):
+    clean_print(LogType.SUCCESS, msg)
+
+
+def log_warning(msg):
+    clean_print(LogType.WARNING, msg)
+
+
+def log_info(msg):
+    clean_print(LogType.INFO, msg)
+
+
+def list_table(list, padding=25, num_cols=4):
+    lines_per_column = int(math.ceil(len(list) / num_cols))
+    for i in range(lines_per_column):
+        for col in range(num_cols):
+            if i + col * lines_per_column < len(list):
+                print(
+                    list[i + col * lines_per_column].ljust(padding),
+                    end="",
+                )
+        print("\n\t", end="")
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD