frontend.py Source File

frontend.py Source File#

Mobilint SDK qb Compiler: frontend.py Source File
Mobilint SDK qb Compiler v0.11.0.1
MCS002-KR
 
import pathlib
from pathlib import Path
from typing import Any
 
import torch
 
from qubee.compiler.compiler import Compiler
from qubee.compiler.compiler import _quantization_task
from qubee.configs import *
from qubee.version import __version__
 
 
 
class Model_Dict(Compiler):
    """
    @brief Wrapper around the Mobilint compiler to support compilation and inference workflows.
 
    @details Capabilities include:
    - Compilation of models into MXQ artifacts runnable on Mobilint NPUs.
    - Inference using the full-precision high-level compiled model on CPU or GPU.
    - Inference with the quantized model on CPU or GPU.
    """
 
    def __init__(
        self,
        model,
        backend="onnx",
        device="cpu",
        feed_dict=None,
        dynamic_axes=None,
        in_dformats=None,
        yolo_decode_include=False,
        **kwargs,
    ):
        """
        @brief Initialize the Mobilint compiler wrapper.
 
        @param model string or model instance. Model path or in-memory model to compile.
        @param backend string. Framework identifier for the model (for example "onnx"). Defaults to "onnx".
        @param device string. Target device for inference ("cpu" or "gpu"). Defaults to "cpu".
        @param feed_dict dict. Example inputs for shape resolution.
        @param dynamic_axes dict. Marks model axes as dynamic.
        @param in_dformats dict. Describes input data formats.
        @param yolo_decode_include bool. Runs YOLO decode on NPU when @c True.
        @param kwargs dict. Additional compiler arguments.
        """
        super().__init__(
            model=model,
            backend=backend,
            device=device,
            feed_dict=feed_dict,
            in_dformats=in_dformats,
            dynamic_axes=dynamic_axes,
            yolo_decode_include=yolo_decode_include,
            **kwargs,
        )
 
    def compile(
        self,
        calib_data_path: str = None,
        save_path: str = None,
        save_subgraph_type: int = 0,
        output_subgraph_path="",
        model_nickname: str = None,
        save_sample=False,
        use_random_calib=False,
        inference_scheme="single",
        cpu_offload=False,
        singlecore_compile=False,
        optimize_option=0,
        sample_dtype: str = "float",
        preprocess_dict=None,
        buffer_mode=1,
        quant_config=QuantConfig.default_config(),
        **kwargs,
    ):
        """
        @brief Compile the wrapped model into an MXQ artifact.
 
        @param calib_data_path string. Calibration dataset directory or meta file (.txt/.json) listing NumPy samples.
        @param save_path string. Output MXQ filename.
        @param save_subgraph_type int. Controls optional MBLT exports: 0 disables; 1 saves graph structure; 2 saves
        structure and weights; 3 splits structure into multiple subgraphs; 4 splits structure and weights.
        @param output_subgraph_path string. Destination for exported .mblt when @p save_subgraph_type is 1–4.
        @param model_nickname string. Friendly name stored in qubee for faster recompile.
        @param save_sample bool. Persists inputs/outputs for debugging in ./sampleInOut.
        @param use_random_calib bool. Uses randomly generated calibration samples when @c True.
        @param inference_scheme string. NPU inference configuration ("single", "multi", "global", etc.).
        @param cpu_offload bool. Enables CPU offload for unsupported groups.
        @param singlecore_compile bool. Forces single-core compilation for large DRAM usage.
        @param optimize_option int. Compiler optimization selector.
        @param sample_dtype string. Data type used for saved sample outputs.
        @param preprocess_dict dict. Additional preprocessing metadata.
        @param buffer_mode int. Buffer serialization mode (0 naive, 1 mmap).
        @param quant_config QuantConfig. Quantization configuration object.
        @param kwargs dict. Additional arguments forwarded to the compiler.
        """
        super().compile(
            calib_data_path=calib_data_path,
            save_path=save_path,
            save_subgraph_type=save_subgraph_type,
            output_subgraph_path=output_subgraph_path,
            model_nickname=model_nickname,
            save_sample=save_sample,
            inference_scheme=inference_scheme,
            use_random_calib=use_random_calib,
            cpu_offload=cpu_offload,
            singlecore_compile=singlecore_compile,
            optimize_option=optimize_option,
            sample_dtype=sample_dtype,
            preprocess_dict=preprocess_dict,
            buffer_mode=buffer_mode,
            quant_config=quant_config,
            **kwargs,
        )
 
 
def mxq_compile(
    model,
    calib_data_path: str = "",
    save_subgraph_type: int = 0,
    output_subgraph_path="",
    model_nickname: str = None,
    save_path: str = None,
    backend="onnx",
    device="cpu",
    feed_dict=None,
    dynamic_axes=None,
    in_dformats=None,
    quantize_method="maxpercentile",
    quantize_percentile=0.9999,
    topk_ratio=0.01,
    smooth_factor=1.6,
    min_clip_ratio=-1,
    act_scale_min=0.0005,
    is_quant_ch=True,
    is_asym_quant=False,
    save_sample=False,
    inference_scheme="single",
    yolo_decode_include=False,
    use_random_calib=False,
    cpu_offload=False,
    singlecore_compile=False,
    optimize_option=0,
    quant_output="layer",
    sample_dtype="float",
    preprocess_dict=None,
    buffer_mode=1,
    weight_dtype="float32",
    use_gpu_only_for_calibration=False,
    input_shape_dict=dict(),
    adaq_config=AdaQuantConfig.default_config(),
    optq_config=OptqConfig.default_config(),
    mod_config=ModConfig.default_config(),
    layer_config=LayerConfig.default_config(),
    llm_config=LlmConfig.default_config(),
    quantbit_config=QuantBitConfig.default_config(),
    **kwargs,
):
    """
    @brief Compile a model into a Mobilint eXeCUtable (MXQ) package for execution on Mobilint NPUs.
 
    @param model string or model instance. Model path. When using @c backend="onnx", this should be the path to an ONNX
    model file. When @c backend is "torchscript", provide a TorchScript module; when "torch", provide a standard
    PyTorch model. For @c backend="tf", pass the directory that contains the TensorFlow SavedModel graph and assets. For
    @c backend="tflite", pass the TF Lite model path.
    @param calib_data_path string. Path to the calibration dataset. Accepts either a text/json file that lists NumPy
    files or a directory that contains the pre-processed NumPy files.
    @param save_subgraph_type int. Controls optional MBLT subgraph exports: 0 disables exports; 1 saves only the graph
    structure; 2 saves graph structure plus weights; 3 saves the graph structure split into multiple subgraphs; 4 saves
    both structure and weights split into multiple subgraphs.
    @param output_subgraph_path string. Destination path for the exported .mblt file when @p save_subgraph_type is 1, 2,
    3, or 4. The resulting file can be used for visualization.
    @param model_nickname string. Friendly name stored in qubee. Defaults to the model basename (for example, the file
    "/workspace/onnx/resnet50.onnx" becomes "resnet50"). Falls back to "temporary" when it cannot be inferred.
    @param save_path string. Output MXQ filename. Defaults to "{model_nickname}.mxq" when omitted.
    @param backend string. Framework used to generate the Mobilint IR. Must be one of "onnx", "tf", "tflite",
    "torchscript", or "torch". Defaults to "onnx".
    @param device string. Compilation and inference device: "cpu" or "gpu". Defaults to "cpu".
    @param feed_dict dict. Example input tensors for shape inference and inference validation.
    @param dynamic_axes dict. Marks model axes as dynamic. Keys are input names and values map dimension indices to
    aliases (for example {"input": {2: "seq_len"}}).
    @param in_dformats dict. Describes input tensor data formats.
    @param quantize_method string. Quantization method. Supported values are "Max", "Percentile", "MaxPercentile",
    "FastPercentile", and "KL". Defaults to "Percentile".
    @param quantize_percentile float. Percentile value for percentile-based quantization; valid range is 0 to 1. Example
    values: 0.999 or 0.9999. Defaults to 0.9999.
    @param topk_ratio float. Applicable to "MaxPercentile" quantization; controls how much data is used for calibration.
    Recommended values are 0.01 or lower; defaults to 0.01.
    @param smooth_factor float. Gaussian kernel smoothing factor required for KL-divergence estimation. Defaults to 1.6.
    @param min_clip_ratio float. Ratio applied to minimum clipping of output scales. Defaults to -1 (disabled).
    @param act_scale_min float. Minimum activation scale used during quantization. Defaults to 0.0005.
    @param is_quant_ch bool. Enables multi-channel quantization when @c True; defaults to @c True.
    @param is_asym_quant bool. Enables asymmetric quantization when @c True; defaults to @c False.
    @param save_sample bool. Stores input/output tensors in ./sampleInOut when @c True to aid debugging. Defaults to
    @c False.
    @param inference_scheme string. NPU inference scheme. One of "single", "multi", "global", "global4", or
    "global8".
    @param yolo_decode_include bool. Determines whether YOLO decode runs on NPU.
    @param use_random_calib bool. Generates random calibration data to validate model compilability. Defaults to @c
    False.
    @param cpu_offload bool. Enables CPU offloading during NPU inference. Defaults to @c False.
    @param singlecore_compile bool. Uses singlecore compilation to maximize DRAM usage. Defaults to @c False.
    @param optimize_option int. Compiler optimization strategy selector.
    @param quant_output string. Quantization applied to the output layer. Accepts "layer" (per-layer), "ch"
    (per-channel; only when @p is_quant_ch is True), or "sigmoid" (use sigmoid-based scaling). Defaults to "layer".
    @param sample_dtype string. Data type for saved sample inference outputs: "float" or "int8".
    @param preprocess_dict int. Identifier for additional pre-processing metadata.
    @param buffer_mode int. Buffer serialization mode: 0 uses a naïve buffer, 1 uses an mmap-backed buffer. Defaults to
    1.
    @param weight_dtype string. Weight dtype for calibration. Defaults to "float32".
    @param use_gpu_only_for_calibration bool. Restricts GPU usage to calibration passes. Defaults to @c False.
    @param input_shape_dict dict. Optional HWC input shape dictionary used for multi-shape compilation (for example
    {"input0": [[224, 224, 3], [256, 256, 3]]}).
    @param adaq_config AdaQuantConfig. Configures the AdaQuant algorithm, including learning rates for weights, bias,
    scales, batch size, and loop epochs. Defaults to disabled values.
    @param optq_config OptqConfig. Configures OPTQ, such as activation-order sorting, block size, and percentile
    dampening.
    @param mod_config ModConfig. Enables Mobilint Optimal Distillation algorithm overrides when provided.
    @param layer_config LayerConfig. Supplies per-layer overrides for enabling MOD/OPTQ, activation bit precision, and
    minimum activation scale values.
    @param llm_config LLMConfig. Large-language-model compilation options such as maximum sequence length, cache length,
    calibration sequence length, mask usage, starting DRAM offset, global-core utilization, batch size, and NPU core IDs.
    @param quantbit_config QuantBitConfig. Controls bit precision for Transformer submodules and optional import/export
    of per-layer precision configurations.
    @param kwargs dict. Additional compiler arguments propagated to the backend implementation.
    @return None.
 
    @par Tips for choosing quantization methods and selecting arguments
    Increase the percentile value to widen the quantization interval. "Percentile" samples the calibration histogram and
    typically starts between 0.9999 and 0.999999. "FastPercentile" relies on reduced, smoothed samples and favors values
    between 0.99 and 0.99999. "MaxPercentile" operates on min/max stacks and commonly uses 0.9 to 0.9999. Channel-wise
    quantization (@p is_quant_ch=True) mitigates per-channel activation variance at the cost of longer compile times.
    Choose @p quant_output="ch" when output variance is channel-dependent; otherwise use "layer".
 
    @par Compiling models with custom inputs (ONNX/Torch/TensorFlow)
    Provide NumPy inputs whenever the model omits shape information so that qubee can infer unknown dimensions and data
    formats.
 
    @code{.py}
    from qubee import mxq_compile
    import numpy as np
 
    example_input = {
        "input_node_name_1": np.random.randn(1, 3, 224, 224).astype(np.float32),
        "input_node_name_2": np.random.randn(1, 8, 224, 224).astype(np.float32),
        "input_node_name_3": np.random.randn(1, 4, 56, 56).astype(np.float32),
    }
 
    in_dformats = {
        "input_node_name_1": "NCHW",
        "input_node_name_2": "NCHW",
        "input_node_name_3": "NCHW",
    }
 
    onnx_model_path = "path/to/your/model.onnx"
    mxq_compile(
        model=onnx_model_path,
        feed_dict=example_input,
        in_dformats=in_dformats,
        backend="onnx",
    )
    @endcode
    """
    if model_nickname and isinstance(model_nickname, str):
        model_name = model_nickname
    elif isinstance(model, str):  # in the case of model path
        model_name = Path(model).stem
    else:  # in the case of model instance
        if save_path is None:
            raise ValueError("Please set the save_path with .mxq extension")
        model_name = "model"
    if save_path is None:
        save_path = f"{model_name}.mxq"
 
    quant_config = get_quant_config(
        quantize_percentile=quantize_percentile,
        topk_ratio=topk_ratio,
        smooth_factor=smooth_factor,
        quantize_method=quantize_method,
        quant_output=quant_output,
        is_quant_ch=is_quant_ch,
        is_asym_quant=is_asym_quant,
        weight_dtype=weight_dtype,
        use_gpu_only_for_calibration=use_gpu_only_for_calibration,
        input_shape_dict=input_shape_dict,
        adaq_config=adaq_config,
        optq_config=optq_config,
        mod_config=mod_config,
        layer_config=layer_config,
        llm_config=llm_config,
        quantbit_config=quantbit_config,
        min_clip_ratio=min_clip_ratio,
        act_scale_min=act_scale_min,
    )
    if isinstance(model, str) and model.endswith(".mblt") and os.path.isfile(model):
        compile_config = get_compile_config(
            calib_data_path=calib_data_path,
            use_random_calib=use_random_calib,
            save_path=save_path,
            save_msgpack_name=None,
            model_nickname=model_name,
            inference_scheme=inference_scheme,
            save_sample=save_sample,
            cpu_offload=cpu_offload,
            singlecore_compile=singlecore_compile,
            optimize_option=optimize_option,
            sample_dtype=sample_dtype,
            preprocess_dict=preprocess_dict,
            version=__version__,
            quant_config=quant_config,
        )
        _quantization_task(model, device, compile_config=compile_config)
    else:
        model_dict = Model_Dict(
            model=model,
            backend=backend,
            device=device,
            feed_dict=feed_dict,
            in_dformats=in_dformats,
            dynamic_axes=dynamic_axes,
            yolo_decode_include=yolo_decode_include,
            **kwargs,
        )
        model_dict.compile(
            model_nickname=model_name,
            calib_data_path=calib_data_path,
            save_path=save_path,
            save_subgraph_type=save_subgraph_type,
            output_subgraph_path=output_subgraph_path,
            save_sample=save_sample,
            inference_scheme=inference_scheme,
            use_random_calib=use_random_calib,
            cpu_offload=cpu_offload,
            singlecore_compile=singlecore_compile,
            optimize_option=optimize_option,
            sample_dtype=sample_dtype,
            preprocess_dict=preprocess_dict,
            buffer_mode=buffer_mode,
            quant_config=quant_config,
        )
        torch.cuda.empty_cache()
        del model_dict
 
 
def mblt_compile(
    model: str | Any,
    mblt_save_path: str,
    backend="onnx",
    device="cpu",
    feed_dict=None,
    dynamic_axes=None,
    in_dformats=None,
    yolo_decode_include=False,
    cpu_offload=False,
    **kwargs,
):
    """
    @brief Export a model to the Mobilint .mblt format without producing an MXQ package.
 
    @param model string or model instance. Source model or path to compile.
    @param mblt_save_path string. Output path for the .mblt artifact.
    @param backend string. Framework identifier ("onnx", "tf", "tflite", "torchscript", or "torch"). Defaults to
    "onnx".
    @param device string. Compilation device ("cpu" or "gpu"). Defaults to "cpu".
    @param feed_dict dict. Example inputs used for shape inference.
    @param dynamic_axes dict. Declares dynamic axes per input name.
    @param in_dformats dict. Input dataformat metadata.
    @param yolo_decode_include bool. Runs YOLO decode on NPU when True.
    @param cpu_offload bool. Enables CPU offloading for unsupported groups.
    @param kwargs dict. Additional arguments forwarded to the compiler.
    @return None.
    """
    model_dict = Model_Dict(
        model=model,
        backend=backend,
        device=device,
        feed_dict=feed_dict,
        dynamic_axes=dynamic_axes,
        in_dformats=in_dformats,
        yolo_decode_include=yolo_decode_include,
        **kwargs,
    )
 
    model_dict.mblt_compile(
        model,
        mblt_save_path=mblt_save_path,
        backend=backend,
        cpu_offload=cpu_offload,
        **kwargs,
    )
 
 
def run_mblt_model(
    path_to_mblt: str,
    save_all_outputs: bool = True,
    inference_output_path: Optional[str] = None,
    target_subgraph: Optional[int] = None,
):
    from qubee.model_dict.common.model_dict import ValueDict
    from qubee.model_dict.inference.model import ModelInference
    from qubee.model_dict.serialize import load_mblt_model
 
    assert path_to_mblt.endswith(".mblt")
    if not save_all_outputs:
        raise NotImplementedError(
            f"save_all_outputs={save_all_outputs} is not supported yet."
        )
    if target_subgraph:
        raise NotImplementedError(
            f"target_subgraph={target_subgraph} is not supported yet."
        )
    import numpy
 
    if not inference_output_path:
        inference_output_path = str(pathlib.Path(path_to_mblt).with_suffix(".infer"))
 
    vd = ValueDict()
    md, wd = load_mblt_model(path_to_mblt, load_weight=True)
    for sg in md.subgraphs:
        for iidx in sg.inputs:
            inact = sg.activations[iidx]
            if inact.name in set(md.inputs):
                vd[inact.name] = numpy.random.random(tuple(map(int, inact.src_shape)))
 
    model_inference = ModelInference(
        model_dict=md,
        weight_dict=wd,
        value_dict=vd,
        save_all_outputs=save_all_outputs,
        inference_all_outputs_write_path=inference_output_path,
    )
 
    model_inference.run_inference(use_value_dict_data_for_next_input=False)
 
 