Source code for aimet_torch.adaround.adaround_weight

# -*- mode: python -*-
# =============================================================================
#  @@-COPYRIGHT-START-@@
#
#  Copyright (c) 2021-2023, Qualcomm Innovation Center, Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  1. Redistributions of source code must retain the above copyright notice,
#     this list of conditions and the following disclaimer.
#
#  2. Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions and the following disclaimer in the documentation
#     and/or other materials provided with the distribution.
#
#  3. Neither the name of the copyright holder nor the names of its contributors
#     may be used to endorse or promote products derived from this software
#     without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
#
#  SPDX-License-Identifier: BSD-3-Clause
#
#  @@-COPYRIGHT-END-@@
# =============================================================================

""" Top level API for Adaptive Rounding - Post-Training Quantization (PTQ) """

import os
import contextlib
import itertools
import json
import tempfile
from typing import Tuple, Union, Dict, List, Callable, Any, Optional
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

# Import AIMET specific modules
from aimet_common.utils import AimetLogger, convert_configs_values_to_bool
from aimet_common.defs import QuantScheme, QuantizationDataType

from aimet_torch import utils
from aimet_torch.save_utils import SaveUtils
from aimet_torch.meta import connectedgraph_utils
from aimet_torch.quantsim import QuantizationSimModel, QcQuantizeWrapper, ExportableQuantModule
from aimet_torch.qc_quantize_op import StaticGridQuantWrapper, QcQuantizeOpMode
from aimet_torch.tensor_quantizer import TensorQuantizer
from aimet_torch.adaround.adaround_wrapper import AdaroundWrapper
from aimet_torch.adaround.adaround_optimizer import AdaroundOptimizer
from aimet_torch.adaround.adaround_loss import AdaroundHyperParameters
from aimet_torch.adaround.activation_sampler import create_modulelist_for_group_modules, get_block_inputs, \
    get_block_outputs, create_cached_block_schedule_list
from aimet_torch.utils import get_named_module

logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Quant)

# The following modules with weights are supported by Adaround
AdaroundSupportedModules = (torch.nn.Conv2d, torch.nn.ConvTranspose2d, torch.nn.Linear)


[docs]class AdaroundParameters: """ Configuration parameters for Adaround """ def __init__(self, data_loader: DataLoader, num_batches: int, default_num_iterations: int = None, default_reg_param: float = 0.01, default_beta_range: Tuple = (20, 2), default_warm_start: float = 0.2, forward_fn: Callable[[torch.nn.Module, Any], Any] = None): """ :param data_loader: Data loader :param num_batches: Number of batches to be used for Adaround. A commonly recommended value for this parameter is the smaller value among (1) len(data_loader) and (2) ceil(2000/batch_size) :param default_num_iterations: Number of iterations to adaround each layer. The default value is 10K for models with 8- or higher bit weights, and 15K for models with lower than 8 bit weights. :param default_reg_param: Regularization parameter, trading off between rounding loss vs reconstruction loss. Default 0.01 :param default_beta_range: Start and stop beta parameter for annealing of rounding loss (start_beta, end_beta). Default (20, 2) :param default_warm_start: warm up period, during which rounding loss has zero effect. Default 20% (0.2) :param forward_fn: Optional adapter function that performs forward pass given a model and inputs yielded from the data loader. The function expects model as first argument and inputs to model as second argument. """ if len(data_loader) < num_batches: raise ValueError(f'Can not fetch {num_batches} batches from ' f'a data loader of length {len(data_loader)}.') self.data_loader = data_loader self.num_batches = num_batches self.num_iterations = default_num_iterations self.reg_param = default_reg_param self.beta_range = default_beta_range self.warm_start = default_warm_start self.forward_fn = forward_fn
class Adaround: """ Weight-rounding mechanism for Post Training Quantization (PTQ) """ @classmethod def apply_adaround(cls, model: torch.nn.Module, dummy_input: Union[torch.Tensor, Tuple], params: AdaroundParameters, path: str, filename_prefix: str, default_param_bw: int = 4, param_bw_override_list: List[Tuple[torch.nn.Module, int]] = None, ignore_quant_ops_list: List[torch.nn.Module] = None, default_quant_scheme: QuantScheme = QuantScheme.post_training_tf_enhanced, default_config_file: str = None) -> torch.nn.Module: """ Returns model with optimized weight rounding of every module (Conv and Linear) and also saves the corresponding quantization encodings to a separate JSON-formatted file that can then be imported by QuantSim for inference or QAT :param model: Model to Adaround :param dummy_input: Dummy input to the model. Used to parse model graph. If the model has more than one input, pass a tuple. User is expected to place the tensors on the appropriate device. :param params: Parameters for Adaround :param path: path where to store parameter encodings :param filename_prefix: Prefix to use for filename of the encodings file :param default_param_bw: Default bitwidth (4-31) to use for quantizing layer parameters :param param_bw_override_list: List of Tuples. Each Tuple is a module and the corresponding parameter bitwidth to be used for that module. :param ignore_quant_ops_list: Ops listed here are skipped during quantization needed for AdaRounding. Do not specify Conv and Linear modules in this list. Doing so, will affect accuracy. :param default_quant_scheme: Quantization scheme. Supported options are using Quant Scheme Enum QuantScheme.post_training_tf or QuantScheme.post_training_tf_enhanced :param default_config_file: Default configuration file for model quantizers :return: Model with Adarounded weights and saves corresponding parameter encodings JSON file at provided path """ # pylint: disable=too-many-arguments # Create Quant sim with given parameters quant_sim = cls._get_quantsim(model, dummy_input=dummy_input, quant_scheme=default_quant_scheme, default_param_bw=default_param_bw, config_file=default_config_file) # For the modules in the param_bw_override_list, override the default parameter bitwidths in the QuantSim if param_bw_override_list: cls._override_param_bitwidth(model, quant_sim, param_bw_override_list) if ignore_quant_ops_list: cls._exclude_modules(model, quant_sim, ignore_quant_ops_list) # Compute only param encodings cls._compute_param_encodings(quant_sim) return cls._apply_adaround(quant_sim, model, dummy_input, params, path, filename_prefix) @classmethod def _apply_adaround(cls, quant_sim: QuantizationSimModel, model: torch.nn.Module, dummy_input: Union[torch.Tensor, Tuple], params: AdaroundParameters, path: str, filename_prefix: str, checkpoints_config: str = None) -> torch.nn.Module: """ Returns model with optimized weight rounding of every module (Conv and Linear) and also saves the corresponding quantization encodings to a separate JSON-formatted file that can then be imported by QuantSim for inference or QAT :param quant_sim: QuantizationSimModel object to optimize weight rounding. The activation quantizers are expected to have been disabled. :param model: Original fp32 model from which quant_sim was created. :param dummy_input: Dummy input to the model. Used to parse model graph. If the model has more than one input, pass a tuple. User is expected to place the tensors on the appropriate device. :param params: Parameters for Adaround :param path: path where to store parameter encodings :param filename_prefix: Prefix to use for filename of the encodings file :param checkpoints_config: Config files to split fp32/quant model by checkpoints :return: Model with Adarounded weights and saves corresponding parameter encodings JSON file at provided path """ # Sanity check: All the input/output quantizers should be disabled cls._check_input_output_quantizers_for_adaround(quant_sim.model) # Get the module - activation function pair using ConnectedGraph module_act_func_pair = connectedgraph_utils.get_module_act_func_pair(model, dummy_input) cls._adaround_model(model, quant_sim, module_act_func_pair, params, dummy_input, checkpoints_config) # Export quantization encodings to JSON-formatted file cls._export_encodings_to_json(path, filename_prefix, quant_sim) cls._remove_quantization_wrappers(quant_sim.model) logger.info('Completed Adarounding Model') return quant_sim.model @classmethod def _adaround_model(cls, model: torch.nn.Module, quant_sim: QuantizationSimModel, module_act_func_pair: Dict, params: AdaroundParameters, dummy_input: Union[torch.Tensor, Tuple], checkpoints_config: str = None): """ Optimize weight rounding of every module (AdaroundSupportedModules) of model in sequential manner based on occurrence NOTE: When checkpoints_config file is provided, assumption is that the outputs from previous group modules (block) should feed directly into next group modules (block) :param model: Original fp32 model from which quant_sim was created. :param quant_sim: QuantizationSimModel object to optimize weight rounding. The activation quantizers are expected to have been disabled. :param module_act_func_pair: Dictionary of module to immediate following activation function :param params: Adaround parameters :param dummy_input: Dummy input to the model :param checkpoints_config: Config files to split fp32/quant model by checkpoints to speedup activations sampling """ # pylint: disable=too-many-locals, protected-access, too-many-branches, too-many-statements num_iterations = params.num_iterations if num_iterations is None: lowest_weight_bw = cls._get_lowest_weight_bw(quant_sim.model) # If the lowest wegith bitwidth is < 8, then set num_iterations to 15K by default if lowest_weight_bw < 8: num_iterations = 15000 else: num_iterations = 10000 with tempfile.TemporaryDirectory() as tmp_dir: # Cache model input data to temporary directory cached_dataset = utils.CachedDataset(params.data_loader, params.num_batches, tmp_dir) # Optimization Hyper parameters opt_params = AdaroundHyperParameters(num_iterations, params.reg_param, params.beta_range, params.warm_start) # AdaRound must be applied to modules in the order of occurrence if checkpoints_config: # Load the predefined json file for checkpoints info checkpoint_config = json.load(open(checkpoints_config)) convert_configs_values_to_bool(checkpoint_config) assert 'cache_on_cpu' in checkpoint_config.keys(), \ "Please define cache_on_cpu to determine whether to cache intermediate tensors on CPU" cache_on_cpu = checkpoint_config['cache_on_cpu'] checkpoint_type = checkpoint_config.get('checkpoint_type', 'sequential') if checkpoint_type == 'sequential': assert 'grouped_modules' in checkpoint_config.keys(), \ "Please provide a dictionary of grouped_modules in the file to define checkpoints" assert 'include_static_inputs' in checkpoint_config.keys(), \ "Please provide a dictionary of include_static_inputs in the file to define checkpoints" grouped_modules = checkpoint_config['grouped_modules'] breakpoint_module_name = checkpoint_config['grouped_modules'][list(grouped_modules.keys())[0]][0] include_static_inputs = checkpoint_config['include_static_inputs'] cached_fp_dataset, cached_quant_dataset = get_block_inputs(model, quant_sim, breakpoint_module_name, cached_dataset, cache_on_cpu, params.forward_fn, params.num_batches, tmp_dir) # Get the device of model to latter be used to place input tensor on the same device device = utils.get_device(model) model.cpu() quant_sim.model.cpu() # Forward function for the ModuleList object def fwd_mod_ls(mod_ls, x): for mod in mod_ls: x = params.forward_fn(mod, x) return x sub_fp_models, sub_sim_models = create_modulelist_for_group_modules(model, quant_sim, grouped_modules) for i, (fp_block, quant_sim_block, static_input) in enumerate(zip(sub_fp_models, sub_sim_models, include_static_inputs)): modules = utils.get_ordered_list_of_modules(fp_block, cached_fp_dataset[0], fwd_mod_ls) cls._run_adaround_model(modules, fp_block, quant_sim_block, module_act_func_pair, opt_params, fwd_mod_ls, cached_fp_dataset, cached_quant_dataset) # Get the outputs from the current block and assign to be the inputs for next block # except for the last block if i < len(sub_fp_models) - 1: get_block_outputs(fp_block, quant_sim_block, static_input, cached_fp_dataset, cached_quant_dataset, cache_on_cpu, fwd_mod_ls, device, tmp_dir) # After finishing Adaround, placing the quant model back to its original device quant_sim.model.to(device) else: assert 'cached_blocks' in checkpoint_config.keys(), \ "Please provide a list of modules that can be cached" block_list = create_cached_block_schedule_list( model, dummy_input, checkpoint_config['cached_blocks'], AdaroundSupportedModules) for block_cfg, modules in tqdm(block_list, desc='block'): if block_cfg is None: # doesn't belong to a cached block cls._run_adaround_model(modules, model, quant_sim.model, module_act_func_pair, opt_params, params.forward_fn, cached_dataset) else: block_name, fp_block = block_cfg quant_sim_block: torch.nn.Module = get_named_module(quant_sim.model, block_name) cached_fp_dataset, cached_quant_dataset = get_block_inputs(model, quant_sim, block_name, cached_dataset, cache_on_cpu, params.forward_fn, params.num_batches, tmp_dir, incl_kwargs=True) def block_fwd(_model, x): return _model(*x) cls._run_adaround_model(modules, fp_block, quant_sim_block, module_act_func_pair, opt_params, block_fwd, cached_fp_dataset, cached_quant_dataset) del cached_fp_dataset del cached_quant_dataset else: modules = utils.get_ordered_list_of_modules(model, dummy_input) cls._run_adaround_model(modules, model, quant_sim.model, module_act_func_pair, opt_params, params.forward_fn, cached_dataset) @classmethod def _run_adaround_model(cls, modules: List, model: torch.nn.Module, quant_sim_model: torch.nn.Module, module_act_func_pair: Dict, opt_params: AdaroundHyperParameters, forward_fn: Callable, cached_dataset: utils.CachedDataset, cached_quant_dataset: Optional[utils.CachedDataset] = None): """ Iterate through all modules to find out Adaround supported modules and apply Adaround optimization to those modules :param modules: Candidate modules :param model: Original fp32 model :param quant_sim_model: QuantSim model :param module_act_func_pair: Activation function pairs :param opt_params: Optimization parameters :param forward_fn: Adapter function that performs forward pass given a model and inputs yielded from the data loader :param cached_dataset: Cached dataset for the fp32 model :param cached_quant_dataset: Cached dataset for the quant model """ # pylint: disable=too-many-arguments, too-many-locals, protected-access for name, module in tqdm(modules): if isinstance(module, AdaroundSupportedModules): # Using name, get corresponding quantized wrapper module from Quant sim model quant_wrapper = cls._get_quant_wrapper(quant_sim_model, name) if not quant_wrapper: continue # Wraps the quant module with adaround wrapper # and temporarily replace quant module with wrapped module with cls._replace_quantization_layer(quant_sim_model, name) as adaround_wrapper: # Get module's next following activation function act_func = module_act_func_pair[module] logger.info("Started Optimizing weight rounding of module: %s", name) AdaroundOptimizer.adaround_module(module, adaround_wrapper, model, quant_sim_model, act_func, cached_dataset, forward_fn, opt_params, cached_quant_dataset) weight = adaround_wrapper.weight # Fold trained alpha to weight with torch.no_grad(): # Use soft rounding to compute Adarounded weight adaround_wrapper.use_soft_rounding = True adarounded_weight = adaround_wrapper.apply_adaround(weight) weight.copy_(adarounded_weight) del adarounded_weight @staticmethod def _compute_param_encodings(quant_sim: QuantizationSimModel): """ Compute encodings for parameters, needed for initializing Adaround quantizers :param quant_sim: Quant sim """ for quant_module in quant_sim.model.modules(): if isinstance(quant_module, StaticGridQuantWrapper): # Adaround requires input and output quantizers to be disabled for quatizer in quant_module.input_quantizers: quatizer.enabled = False for quatizer in quant_module.output_quantizers: quatizer.enabled = False # pylint: disable=protected-access for name, param in quant_module._module_to_wrap.named_parameters(): param_quantizer = quant_module.param_quantizers[name] param_quantizer.reset_encoding_stats() param_quantizer.update_encoding_stats(param.data) param_quantizer.compute_encoding() # Wrapper mode must be set to ACTIVE because the wrapper's quantize_dequantize_params() will only call # into the param tensor quantizer's quantize_dequantize() if the mode is not PASSTHROUGH. quant_module.set_mode(QcQuantizeOpMode.ACTIVE) @staticmethod def _get_quantsim(model: torch.nn.Module, dummy_input: torch.Tensor, quant_scheme: QuantScheme, default_param_bw: int, config_file: str): return QuantizationSimModel(model, dummy_input=dummy_input, quant_scheme=quant_scheme, default_param_bw=default_param_bw, config_file=config_file) @staticmethod def _get_adaround_wrapper(quant_module: QcQuantizeWrapper): return AdaroundWrapper(quant_module) @staticmethod def _remove_quantization_wrappers(module: torch.nn.Module): SaveUtils.remove_quantization_wrappers(module) @staticmethod @contextlib.contextmanager def _patch_module_layer(model, layer_name, new_layer): """ Temporarily replace model layer """ original_layer = getattr(model, layer_name) setattr(model, layer_name, new_layer) yield setattr(model, layer_name, original_layer) @staticmethod def _validate_quant_module_for_adaround(quant_module: StaticGridQuantWrapper): assert quant_module.param_quantizers['weight'], '%s does not have weight parameter.' % quant_module assert quant_module.param_quantizers['weight'].encoding, '%s encoding needs to be set.' % quant_module @staticmethod def _check_input_output_quantizers_for_adaround(quant_model: torch.nn.Module): _, input_quantizers, output_quantizers = utils.get_all_quantizers(quant_model) for quantizer in itertools.chain(input_quantizers, output_quantizers): assert not quantizer.enabled @staticmethod def _get_lowest_weight_bw(quant_model: torch.nn.Module): param_quantizers, _, _ = utils.get_all_quantizers(quant_model) return min( quantizer.bitwidth for quantizer in param_quantizers if quantizer.enabled and quantizer.data_type == QuantizationDataType.int ) @classmethod @contextlib.contextmanager def _replace_quantization_layer(cls, quant_sim_model: torch.nn.Module, module_name: str): """ Replace the quantized module's weight tensor quantizer with the Adaround tensor quantizer :param quant_module: quant module """ quant_module = utils.get_named_module(quant_sim_model, module_name) cls._validate_quant_module_for_adaround(quant_module) adaround_layer = cls._get_adaround_wrapper(quant_module) # We need to look for the container to patch for modules inside submodule upper_module = quant_sim_model upper_module_name, _, target_module_name = module_name.rpartition('.') if upper_module_name: upper_module = utils.get_named_module(quant_sim_model, upper_module_name) # Temporarily replace quant module with wrapped module with cls._patch_module_layer(upper_module, target_module_name, adaround_layer): yield adaround_layer @staticmethod def _get_quant_wrapper(quant_sim_model: torch.nn.Module, module_name: str) -> Union[StaticGridQuantWrapper, None]: """ For given module name, get the quantized wrapper module from the QuantSim model :param quant_sim_model: Model with simulation ops :param module_name: Module name :return: Quantized wrapper module or None """ quant_module = None for name, module in quant_sim_model.named_modules(): if name == module_name and isinstance(module, StaticGridQuantWrapper): quant_module = module break return quant_module @classmethod def _export_encodings_to_json(cls, path: str, filename_prefix: str, quant_sim: QuantizationSimModel): """ Save Adadrounded module's parameter encodings to JSON file :param path: path where to store param encodings :param filename_prefix: filename to store exported weight encodings in JSON format :param quant_sim: QunatSim that contains the model and Adaround tensor quantizers """ # pylint: disable=protected-access # Create a dictionary to export to JSON file param_encodings = {} for name, quant_module in quant_sim.model.named_modules(): if isinstance(quant_module, ExportableQuantModule) and \ isinstance(quant_module.get_original_module(), AdaroundSupportedModules): if 'weight' in quant_module.param_quantizers: cls._update_param_encodings_dict(quant_module, name, param_encodings) # Unify the encoding format to be same as that of full encoding export file encoding = {'param_encodings': param_encodings} # export encodings to JSON file os.makedirs(os.path.abspath(path), exist_ok=True) encoding_file_path = os.path.join(path, filename_prefix + '.encodings') with open(encoding_file_path, 'w') as encoding_fp: json.dump(encoding, encoding_fp, sort_keys=True, indent=4) @classmethod def _update_param_encodings_dict(cls, quant_module: ExportableQuantModule, name: str, param_encodings: Dict): """ Add module's weight parameter encodings to dictionary to be used for exporting encodings :param quant_module: quant module :param name: name of module :param param_encodings: Dictionary of param encodings """ for orig_param_name, encodings in quant_module.export_param_encodings().items(): if orig_param_name == 'weight' and encodings: param_name = name + '.' + orig_param_name param_encodings[param_name] = encodings @staticmethod def _create_encodings_dict_for_quantizer(quantizer: TensorQuantizer) -> List[Dict]: """ Return encodings for given qunatizer :param quantizer: Tensor quantizer associated with module's param :return: Dictionary containing encodings """ quant_encodings = quantizer.encoding if not isinstance(quantizer.encoding, list): quant_encodings = [quant_encodings] encodings_dict = [] for enc in quant_encodings: encodings_dict.append({'min': enc.min, 'max': enc.max, 'scale': enc.delta, 'offset': int(enc.offset), 'bitwidth': enc.bw, 'is_symmetric': str(quantizer.use_symmetric_encodings), 'dtype': 'int' if quantizer.data_type == QuantizationDataType.int else 'float'}) return encodings_dict @staticmethod def _override_param_bitwidth(model: torch.nn.Module, quant_sim: QuantizationSimModel, param_bw_override_list: List[Tuple[torch.nn.Module, int]]): """ For the QuantSim, for the list of modules in the param_bw_override_list, overrides the default parameter bitwidths with the provided bitwidth. :param model: The original model :param quant_sim: The QuantSim that was created using a deepcopy of the original model. :param param_bw_override_list: List of Tuples. Each Tuple is a module and the corresponding parameter bitwidth to be used for that module. """ # Create a mapping of original model's AdaRoundable module and their name module_to_name = {} for name, module in model.named_modules(): if isinstance(module, AdaroundSupportedModules): module_to_name[module] = name # Create a mapping of QuantSim model's AdaRoundable module name and their module name_to_module = {} for q_name, q_module in quant_sim.model.named_modules(): if isinstance(q_module, ExportableQuantModule): if isinstance(q_module.get_original_module(), AdaroundSupportedModules): # pylint: disable=protected-access name_to_module[q_name] = q_module # For the modules specified in the param_bw_override_list, set the weight quantizer bitwidth for (module, bw) in param_bw_override_list: module_name = module_to_name[module] quant_wrapper = name_to_module[module_name] quant_wrapper.param_quantizers['weight'].bitwidth = bw @classmethod def _exclude_modules(cls, model: torch.nn.Module, quant_sim: QuantizationSimModel, ignore_quant_ops_list: List[torch.nn.Module]): """ For the modules mentioned in the ignore_quant_ops_list, remove the corresponding quant wrappers from the quantSim and excludes modules from adaround optimization. :param model: The original model :param quant_sim: The QuantSim that was created using a deepcopy of the original model. :param ignore_quant_ops_list: The list of modules for which the Quantization wrappers are removed from the QuantSim object. """ quant_wrappers_to_exclude = [] for module in ignore_quant_ops_list: for m in module.modules(): name = utils.get_layer_name(model, m) quant_wrapper = cls._get_quant_wrapper(quant_sim.model, name) if quant_wrapper: quant_wrappers_to_exclude.append(quant_wrapper) quant_sim.exclude_layers_from_quantization(quant_wrappers_to_exclude) @classmethod def apply_adaround_with_cache(cls, model: torch.nn.Module, dummy_input: Union[torch.Tensor, Tuple], params: AdaroundParameters, path: str, filename_prefix: str, default_param_bw: int = 4, param_bw_override_list: List[Tuple[torch.nn.Module, int]] = None, ignore_quant_ops_list: List[torch.nn.Module] = None, default_quant_scheme: QuantScheme = QuantScheme.post_training_tf_enhanced, default_config_file: str = None, checkpoints_config: str = None) -> torch.nn.Module: """ Returns model with optimized weight rounding of every module (Conv and Linear) and also saves the corresponding quantization encodings to a separate JSON-formatted file that can then be imported by QuantSim for inference or QAT :param model: Model to Adaround :param dummy_input: Dummy input to the model. Used to parse model graph. If the model has more than one input, pass a tuple. User is expected to place the tensors on the appropriate device. :param params: Parameters for Adaround :param path: path where to store parameter encodings :param filename_prefix: Prefix to use for filename of the encodings file :param default_param_bw: Default bitwidth (4-31) to use for quantizing layer parameters :param param_bw_override_list: List of Tuples. Each Tuple is a module and the corresponding parameter bitwidth to be used for that module. :param ignore_quant_ops_list: Ops listed here are skipped during quantization needed for AdaRounding. Do not specify Conv and Linear modules in this list. Doing so, will affect accuracy. :param default_quant_scheme: Quantization scheme. Supported options are using Quant Scheme Enum QuantScheme.post_training_tf or QuantScheme.post_training_tf_enhanced :param default_config_file: Default configuration file for model quantizers :param checkpoints_file: JSON file to define checkpoints for caching intermediate tensors of fp32/quant model :return: Model with Adarounded weights and saves corresponding parameter encodings JSON file at provided path """ # pylint: disable=too-many-arguments assert checkpoints_config is not None, "To run Adaround with cached tensors, please provide a JSON file with checkpoints defined" # Create Quant sim with given parameters quant_sim = cls._get_quantsim(model, dummy_input=dummy_input, quant_scheme=default_quant_scheme, default_param_bw=default_param_bw, config_file=default_config_file) # For the modules in the param_bw_override_list, override the default parameter bitwidths in the QuantSim if param_bw_override_list: cls._override_param_bitwidth(model, quant_sim, param_bw_override_list) if ignore_quant_ops_list: cls._exclude_modules(model, quant_sim, ignore_quant_ops_list) # Compute only param encodings cls._compute_param_encodings(quant_sim) return cls._apply_adaround(quant_sim, model, dummy_input, params, path, filename_prefix, checkpoints_config)