compile_config.py Source File

compile_config.py Source File#

Mobilint SDK qb Compiler: compile_config.py Source File
Mobilint SDK qb Compiler v0.11.0.1
MCS002-EN
from typeguard import typechecked
from . import ConfigABC, QuantConfig, is_valid_dir_path, check_save_path
import tempfile
import numbers
from qubee.calibration.utils_calib import (
    list_np_files_in_json,
    list_np_files_in_txt,
    CalibType,
    check_calib_data,
)
from qubee.version import __version__
 
 
class CompileConfig(ConfigABC):
    py_to_cpp_name = {
        "calib_data_path": "calibData",
        "save_path": "savePath",
        "model_nickname": "modelNickname",
        "save_msgpack_name": "saveMsgpackName",
        "model_path": "modelPath",
        "inference_scheme": "inferenceScheme",
        "save_sample": "saveSample",
        "cpu_offload": "cpuOffload",
        "singlecore_compile": "singlecoreCompile",
        "optimize_option": "optimizeOption",
        "sample_dtype": "sampleDtype",
        "preprocess_dict": "preprocess",
        "version": "version",
    }
    min_max_check = {"optimize_option": [0, None]}
    inference_scheme_types = ["single", "multi", "global", "global4", "global8"]
    sample_dtypes = ["float", "int8"]
 
    @typechecked
    def __init__(
        self,
        calib_data_path: str | None,
        use_random_calib: bool,
        save_path: str,
        model_nickname: str,
        save_msgpack_name: str | None,
        model_path: str | None,
        inference_scheme: str,
        save_sample: bool,
        cpu_offload: bool,
        singlecore_compile: bool,
        optimize_option: int,
        sample_dtype: str,
        preprocess_dict: dict | None,
        buffer_mode: int,
        version: str,
        quant_config: QuantConfig,
        optional: bool = False,
    ):
        super().__init__(optional)
        self.calib_data_path = self._get_calibration_path(
            calib_data_path, use_random_calib
        )
        self.save_path = save_path
        self.model_nickname = model_nickname
        self.save_msgpack_name = save_msgpack_name
        self.model_path = model_path
        self.inference_scheme = inference_scheme
        self.save_sample = save_sample
        self.cpu_offload = cpu_offload
        self.singlecore_compile = singlecore_compile
        self.optimize_option = optimize_option
        self.sample_dtype = sample_dtype
        self.preprocess_dict = preprocess_dict
        self.version = version
        self.quant_config = quant_config
        self.buffer_mode = buffer_mode
        self.check_valid(
            except_list=["optional", "quant_config", "calib_file", "buffer_mode"]
        )  # quant_config is already verified
 
    @classmethod
    def default_config(cls):
        calib_data_path = None
        use_random_calib = False
        save_path = "./tmp.mxq"
        save_msgpack_name = None
        model_nickname = "temporary"
        model_path = None
        inference_scheme = "single"
        save_sample = False
        cpu_offload = False
        singlecore_compile = False
        optimize_option = 0
        sample_dtype = "float"
        preprocess_dict = None
        buffer_mode = 1
        quant_config = QuantConfig.default_config()
        version = __version__
 
        return CompileConfig(
            optional=False,
            calib_data_path=calib_data_path,
            use_random_calib=use_random_calib,
            save_path=save_path,
            save_msgpack_name=save_msgpack_name,
            model_nickname=model_nickname,
            model_path=model_path,
            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,
            buffer_mode=buffer_mode,
            version=version,
            quant_config=quant_config,
        )
 
    def check_valid_str_configs(self):
        if self.save_path:
            check_save_path(self.save_path, "mxq")
        if self.save_msgpack_name:
            is_valid_dir_path(self.save_msgpack_name)
        if self.inference_scheme.lower() not in CompileConfig.inference_scheme_types:
            raise ValueError(
                f"Inference_scheme should be one of the {CompileConfig.inference_scheme_types}"
            )
        if self.sample_dtype.lower() not in CompileConfig.sample_dtypes:
            raise ValueError(
                f"Sample data type should be one of the {CompileConfig.sample_dtypes}"
            )
        if self.buffer_mode not in (0, 1):
            raise ValueError(f"buffer mode must be either 0 or 1")
 
    def to_dict(self):
        compile_config_dict = dict()
        for py_name, cpp_name in CompileConfig.py_to_cpp_name.items():
            value = getattr(self, py_name)
            if isinstance(value, str) and py_name in [
                "inference_scheme",
                "sample_dtype",
            ]:
                value = value.lower()
            compile_config_dict[cpp_name] = value
        compile_config_dict["quant"] = self.quant_config.to_dict()
 
        return compile_config_dict
 
    def _get_calibration_path(self, calib_data_path: str, use_random_calib: bool):
        if not calib_data_path and not use_random_calib:
            raise ValueError(
                "Please use calib_data_path or enable use_random_calib. Do not leave calib_data_path empty while setting use_random_calib to false."
            )
        if (
            not calib_data_path or use_random_calib
        ):  # we don't have to check the validity of calib_data_path
            return None
        calib_data_case = check_calib_data(calib_data_path)
        if calib_data_case == CalibType.SINGLE_DIR:
            self.calib_file = tempfile.NamedTemporaryFile(suffix=".txt")
            list_np_files_in_txt(calib_data_path, self.calib_file.name)
            calib_data_path = self.calib_file.name
        elif calib_data_case == CalibType.MULTI_DIR:
            self.calib_file = tempfile.NamedTemporaryFile(suffix=".json")
            list_np_files_in_json(calib_data_path, self.calib_file.name)
            calib_data_path = self.calib_file.name
        elif calib_data_case in (CalibType.SINGLE_TXT, CalibType.MULTI_JSON):
            pass
        else:
            raise ValueError(f"Got unexpected calib_data_path={calib_data_path}.")
 
        return calib_data_path
 
 
def get_compile_config(
    calib_data_path=None,
    use_random_calib=False,
    save_path="./tmp.mxq",
    save_msgpack_name=None,
    model_nickname="temporary",
    model_path=None,
    inference_scheme="single",
    save_sample=False,
    cpu_offload=False,
    singlecore_compile=False,
    optimize_option=0,
    sample_dtype="float",
    preprocess_dict=None,
    buffer_mode=1,
    quant_config=QuantConfig.default_config(),
    version=__version__,
):
    return CompileConfig(
        optional=False,
        calib_data_path=calib_data_path,
        use_random_calib=use_random_calib,
        save_path=save_path,
        save_msgpack_name=save_msgpack_name,
        model_nickname=model_nickname,
        model_path=model_path,
        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,
        buffer_mode=buffer_mode,
        version=version,
        quant_config=quant_config,
    )