Source code for afe.core.calibrate_networks

#########################################################
# 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 Union, Iterable, Dict

import numpy as np

from sima_utils.data.data_generator import DataGenerator

from afe.core.configs import OptimizationConfigs
from afe.driver import passes
from afe.ir.defines import Status, NodeName
from afe.ir.net import AwesomeNet
from afe.core.utils import convert_data_generator_to_iterable


[docs] def calibrate_network(net: AwesomeNet, opt_config: OptimizationConfigs, input_generator: Union[DataGenerator, Iterable[Dict[NodeName, np.ndarray]]]) -> None: """ Save quantization configuration in an AwesomeNet, then calibrate it. This function is used in legacy tests. Do not call this function. :param net: an AwesomeNet :param opt_config: A OptimizationConfigs instance containing quantization scheme information. :param input_generator: Source of input values for calibration. """ # Prepare to run calibration if isinstance(input_generator, DataGenerator): input_generator = convert_data_generator_to_iterable(input_generator) calibration_pass = passes.update_quantization_configs(opt_config.quantization_configs)(net) \ .then(lambda net1: passes.calibration(opt_config.calibration_configs)(net1, input_generator)) # Run pipeline net_before_calibration = net net = calibration_pass.run() assert net.status == Status.CALIBRATED assert net is net_before_calibration # This function promises to modify net in-place