Function quantize_model

neural_compressor/jax/quantization/quantize.py:43–71

Return a quantized Keras model according to the given configuration. Args: model (keras.Model): FP32 Keras model to be quantized. quant_config (BaseConfig): Quantization configuration. calib_function (Callable, optional): Function used for model calibration, required for

(
    model: keras.Model,
    quant_config: BaseConfig,
    calib_function: Callable = None,
    inplace: bool = True
)

# fmt: off
@log_process(mode=Mode.QUANTIZE)
def quantize_model(
    model: keras.Model,
    quant_config: BaseConfig,
    calib_function: Callable = None,
    inplace: bool = True
):
    """Return a quantized Keras model according to the given configuration.

    Args:
        model (keras.Model): FP32 Keras model to be quantized.
        quant_config (BaseConfig): Quantization configuration.
        calib_function (Callable, optional): Function used for model calibration, required for static quantization.
        inplace (bool): When True, the original model is modified in-place and should not be used afterward. False creates a copy of original model

    Returns:
        keras.Model: The quantized model.
    """
# fmt: on
    check_backend()
    if not inplace:
        model = clone_model(model)

    model_info = quant_config.get_model_info(model)
    configs_mapping = quant_config.to_config_mapping(model_info=model_info)
    for algo_name, algo_func in algos_mapping.items():
        if need_apply(configs_mapping, algo_name):
            logger.info(f"Start to apply {algo_name} on the model.")
            model = algo_func(model, configs_mapping, quant_config, calib_function)
    return model