Class QuantizeInputs

exir/passes/quantize_io_pass.py:205–266

class QuantizeInputs(ExportPass):
    def __init__(
        self,
        edge_program_manager: EdgeProgramManager,
        quantized_inputs_idx: Union[Dict[int, Dict[str, Any]], List[int]],
        method_name: Optional[str] = None,
        exported_program: Optional[ExportedProgram] = None,
    ):
        super().__init__()
        self.edge_program_manager = edge_program_manager

        self.quantized_inputs_idx_dict = {}
        if isinstance(quantized_inputs_idx, dict):
            self.quantized_inputs_idx_dict = quantized_inputs_idx
        else:
            for idx in quantized_inputs_idx:
                self.quantized_inputs_idx_dict[idx] = None
        self.param_prefix_name = method_name
        self.exported_program = exported_program
        self.quant_args = {}

    def edge_manager_update_quant_config_method(self, idx, quant_args):
        if self.edge_program_manager is not None:
            if not self.edge_program_manager._config_methods:
                self.edge_program_manager._config_methods = {}

            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", idx, "scale")
            ] = quant_args[0]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", idx, "zp")
            ] = quant_args[1]
            self.edge_program_manager._config_methods[
                get_config_method_name(
                    self.param_prefix_name, "input", idx, "quant_min"
                )
            ] = quant_args[2]
            self.edge_program_manager._config_methods[
                get_config_method_name(
                    self.param_prefix_name, "input", idx, "quant_max"
                )
            ] = quant_args[3]
            self.edge_program_manager._config_methods[
                get_config_method_name(self.param_prefix_name, "input", idx, "dtype")
            ] = scalar_type_enum(quant_args[4])

    def edge_manager_update_quant_config_methods_all(self):
        if self.edge_program_manager is not None:
            for idx, val in self.quant_args.items():
                self.edge_manager_update_quant_config_method(idx, val)

    def call(self, graph_module: torch.fx.GraphModule):
        for i, qparams in self.quantized_inputs_idx_dict.items():
            exported_program = (
                self.edge_program_manager.exported_program()
                if self.edge_program_manager is not None
                else self.exported_program
            )