resource_management_config.py Source File

resource_management_config.py Source File#

Mobilint SDK qb Compiler: resource_management_config.py Source File
Mobilint SDK qb Compiler v0.12.0.0
MCS002-EN
from typeguard import typechecked
from . import ConfigABC
 
 
class ResourceManagementConfig(ConfigABC):
    """
    @brief Configuration for resource management during model compilation.
    """
 
    optional = False
    DEFAULTS = {
        "weight_dtype": "float32",
        "use_gpu_only_for_calibration": True,
        "weight_memory_method": 1,
    }
 
    py_to_cpp_name = {
        "weight_dtype": "weightDtype",
        "use_gpu_only_for_calibration": "useGPUOnlyForCalibration",
        "weight_memory_method": "weightMemoryMethod",
    }
 
    min_max_check = {
        "weight_memory_method": [0, 4],
    }
 
    weight_memory_methods = [
        "DeleteFloat",
        "SaveFloat",
        "MoveFloat",
        "KeepFloat",
        "KeepAll",
    ]
 
    def __init__(
        self,
        weight_dtype: str,
        use_gpu_only_for_calibration: bool,
        weight_memory_method: int,
        optional: bool = optional,
    ):
        """
        @brief Initialize the ResourceManagementConfig.
        @param weight_dtype str. Weight data type for calibration (e.g., "float32", "float16").
        @param use_gpu_only_for_calibration bool. If True, use GPU only during the calibration phase.
        @param weight_memory_method int. Weight memory management method index.<br>
            0: DeleteFloat - Delete float weights after quantization.<br>
            1: SaveFloat - Save float weights to disk.<br>
            2: MoveFloat - Move float weights to CPU.<br>
            3: KeepFloat - Keep float weights in memory.<br>
            4: KeepAll - Keep all weights in memory.<br>
        @param optional bool. Indicates whether this config is optional. Use the default value defined as a class variable.
        """
        super().__init__(optional)
        self.weight_dtype = weight_dtype
        self.use_gpu_only_for_calibration = use_gpu_only_for_calibration
        self.weight_memory_method = weight_memory_method
        self.check_valid()
 
    @classmethod
    def default_config(cls):
        return cls(**cls.DEFAULTS, optional=cls.optional)
 
    @classmethod
    def from_kwargs(cls, **kwargs):
        params = cls.DEFAULTS.copy()
        updated_kwargs = dict()
        for key in cls.DEFAULTS.keys():
            if key in kwargs:
                updated_kwargs[key] = kwargs[key]
                params[key] = kwargs[key]
 
        try:
            return cls(**params, optional=cls.optional)
        except:
            raise ValueError(
                f"Please review the configuration kwargs passed to {cls.__name__}. These kwargs were updated and will be used: {updated_kwargs}"
            )
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        """Create ResourceManagementConfig from dictionary (JSON structure)"""
        params = cls.DEFAULTS.copy()
 
        if "weightDtype" in config_dict:
            params["weight_dtype"] = config_dict["weightDtype"]
        if "useGPUOnlyForCalibration" in config_dict:
            params["use_gpu_only_for_calibration"] = config_dict[
                "useGPUOnlyForCalibration"
            ]
 
        if "weightMemory" in config_dict:
            weight_memory = config_dict["weightMemory"]
            if "method" in weight_memory:
                params["weight_memory_method"] = weight_memory["method"]
 
        return cls(**params)
 
    def to_dict(self):
        return {
            "resourceManagement": {
                "weightDtype": self.weight_dtype,
                "useGPUOnlyForCalibration": self.use_gpu_only_for_calibration,
                "weightMemory": {
                    "methodList": self.weight_memory_methods,
                    "method": self.weight_memory_method,
                },
            }
        }
 
 
def get_resource_management_config(**kwargs) -> ResourceManagementConfig:
    """
    @brief Create ResourceManagementConfig with partial parameters merged with defaults.
    """
    params = ResourceManagementConfig.DEFAULTS.copy()
    params.update(kwargs)
    return ResourceManagementConfig(
        **params, optional=ResourceManagementConfig.optional
    )
 
 
# @}