#########################################################
# Copyright (C) 2020 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Joey Chou
#########################################################
from typing import Optional, Dict, Any, Union, Iterable
import numpy as np
from sima_utils.data.data_generator import DataGenerator
import afe
from afe.core.configs import ModelConfigs, OptimizationConfigs
from afe.driver import passes
from afe.driver.compile_step import CompileStep
from afe.driver.passes import dump_diagnostic_files_after
from afe.ir.defines import Status, NodeName, BiasCorrectionType
from afe.ir.net import AwesomeNet
from afe.core.utils import save_files, convert_data_generator_to_iterable
[docs]
def quantize_network(net: AwesomeNet,
model_config: ModelConfigs,
opt_config: OptimizationConfigs,
input_generator: Union[DataGenerator, Iterable[Dict[NodeName, np.ndarray]], None] = None,
dump_to_files: bool = False,
custom_quantization_configs: Optional[Dict[NodeName, Dict[str, Any]]] = None) -> None:
"""
Quantizes the AwesomeNet. If the AwesomeNet has not been calibrated, run the calibration first.
Save the YAML, JSON, and npz if the SIMA_AFE_SAVED_FILES environmental variables is set to `1`
or dump_to_files input argument has been set to True.
:param net: an AwesomeNet
:param model_config: A ModelConfigs instance containing model related information and status.
:param opt_config: A OptimizationConfigs instance containing quantization scheme information.
:param input_generator: Source of input values for calibration. May be None if network is already calibrated.
:param dump_to_files: Flag enabling writing the quantized AwesomeNet to .npz file and configs to .yaml file
:param custom_quantization_configs: Optional[Dict[NodeName, Dict[str, Any]]]. A dictionary using NodeName
as keys. The value to each key is a dictionary of the AwesomeQuantAttr's field names and sets target configuration.
Example
-------
The example shows how a custom_quantization_configs looks like to config the output_int32 field in a Conv2DQuantAttrs
in a output conv2d_add node to True.
custom_quantization_configs = {"MLA_1/conv2d_add_84": {"output_int32": True}}
"""
set_quantization_configs = passes.update_quantization_configs(
opt_config.quantization_configs,
custom_quantization_configs=custom_quantization_configs
)
equalize = passes.equalization(opt_config.calibration_configs, opt_config.quantization_configs)
calibrate = passes.calibration(opt_config.calibration_configs)
quantize = passes.quantization(
input_generator if opt_config.quantization_configs.biascorr_type.get() == BiasCorrectionType.ITERATIVE
else None
)
def dump_after(s: CompileStep[AwesomeNet], suffix: str) -> CompileStep[AwesomeNet]:
return dump_diagnostic_files_after(s, model_config, opt_config,
condition=save_files() or dump_to_files,
suffix=suffix)
# Prepare to run calibration if needed
if net.status != Status.CALIBRATED:
if input_generator is None:
raise ValueError("The AwesomeNet is not calibrated, please provide dataset input generator.")
if isinstance(input_generator, DataGenerator):
input_generator = convert_data_generator_to_iterable(input_generator)
calibration_step = set_quantization_configs(net) \
.then(lambda net1: equalize(net1, input_generator)) \
.then(lambda net2: calibrate(net2, input_generator))
else:
# Already calibrated. Do nothing.
calibration_step = CompileStep.pure(net)
# Prepare to run quantization
quantization_step = calibration_step \
.then(lambda c_net: dump_after(quantize(c_net), afe.QUANTIZED_POSTFIX))
# Run pipeline
net_before_quantization = net
net = quantization_step.run()
assert net is net_before_quantization # This function promises to modify net in-place