|
| 1 | +#!/usr/bin/env python |
| 2 | +# -*- coding: utf-8 -*- |
| 3 | +# |
| 4 | +# Copyright (c) 2024 Intel Corporation |
| 5 | +# |
| 6 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | +# you may not use this file except in compliance with the License. |
| 8 | +# You may obtain a copy of the License at |
| 9 | +# |
| 10 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | +# |
| 12 | +# Unless required by applicable law or agreed to in writing, software |
| 13 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | +# See the License for the specific language governing permissions and |
| 16 | +# limitations under the License. |
| 17 | + |
| 18 | +import json |
| 19 | + |
| 20 | +from neural_compressor.torch.utils import get_ipex_version |
| 21 | + |
| 22 | +try: |
| 23 | + import intel_extension_for_pytorch as ipex |
| 24 | +except: |
| 25 | + assert False, "Please install IPEX for static quantization." |
| 26 | + |
| 27 | +import torch |
| 28 | +from packaging.version import Version |
| 29 | + |
| 30 | +from .utility import ( |
| 31 | + cfg_to_qconfig, |
| 32 | + dump_model_op_stats, |
| 33 | + get_quantizable_ops_recursively, |
| 34 | + ipex_config_path, |
| 35 | + simple_inference, |
| 36 | +) |
| 37 | + |
| 38 | +ipex_ver = get_ipex_version() |
| 39 | + |
| 40 | + |
| 41 | +def static_quantize(model, tune_cfg, run_fn, example_inputs, inplace=True): |
| 42 | + """Execute the quantize process on the specified model. |
| 43 | +
|
| 44 | + Args: |
| 45 | + model: a float model to be quantized. |
| 46 | + tune_cfg: quantization config for ops. |
| 47 | + run_fn: a calibration function for calibrating the model. |
| 48 | + example_inputs: used to trace torch model. |
| 49 | + inplace: whether to carry out model transformations in-place. |
| 50 | +
|
| 51 | + Returns: |
| 52 | + A quantized model. |
| 53 | + """ |
| 54 | + model.eval() |
| 55 | + |
| 56 | + if ipex_ver.release >= Version("1.12.0").release: |
| 57 | + # Check save_qconf_summary part is a workaround for IPEX bug. |
| 58 | + # Sometimes the prepared model from get_op_capablitiy loss this attribute |
| 59 | + if not hasattr(model, "save_qconf_summary") or not hasattr(model, "load_qconf_summary"): |
| 60 | + from torch.ao.quantization import MinMaxObserver, PerChannelMinMaxObserver, QConfig |
| 61 | + |
| 62 | + if ipex_ver.release >= Version("2.1").release: |
| 63 | + static_qconfig = ipex.quantization.default_static_qconfig_mapping |
| 64 | + else: |
| 65 | + static_qconfig = QConfig( |
| 66 | + activation=MinMaxObserver.with_args(qscheme=torch.per_tensor_affine, dtype=torch.quint8), |
| 67 | + weight=PerChannelMinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_channel_symmetric), |
| 68 | + ) |
| 69 | + if isinstance(example_inputs, dict): |
| 70 | + model = ipex.quantization.prepare( |
| 71 | + model, static_qconfig, example_kwarg_inputs=example_inputs, inplace=inplace |
| 72 | + ) |
| 73 | + else: |
| 74 | + model = ipex.quantization.prepare(model, static_qconfig, example_inputs=example_inputs, inplace=inplace) |
| 75 | + |
| 76 | + model.load_qconf_summary(qconf_summary=ipex_config_path) |
| 77 | + run_fn(model) |
| 78 | + model.save_qconf_summary(qconf_summary=ipex_config_path) |
| 79 | + model = _ipex_post_quant_process(model, example_inputs, inplace=inplace) |
| 80 | + |
| 81 | + else: # pragma: no cover |
| 82 | + # for IPEX version < 1.12 |
| 83 | + _, cfgs, default_cfgs, fuse_ops = get_quantizable_ops_recursively(model, example_inputs) |
| 84 | + qscheme = cfg_to_qconfig(tune_cfg, cfgs, default_cfgs, fuse_ops) |
| 85 | + ipex_conf = ipex.quantization.QuantConf( |
| 86 | + configure_file=ipex_config_path, qscheme=qscheme |
| 87 | + ) # pylint: disable=E1101 |
| 88 | + run_fn(model) |
| 89 | + ipex_conf.save(ipex_config_path) |
| 90 | + ipex_conf = ipex.quantization.QuantConf(ipex_config_path) # pylint: disable=E1101 |
| 91 | + model = ipex.quantization.convert(model, ipex_conf, example_inputs, inplace=True) # pylint: disable=E1121 |
| 92 | + |
| 93 | + with open(ipex_config_path, "r") as f: |
| 94 | + model.tune_cfg = json.load(f) |
| 95 | + model.ipex_config_path = ipex_config_path |
| 96 | + if ipex_ver.release >= Version("1.12.0").release: |
| 97 | + dump_model_op_stats(tune_cfg) |
| 98 | + return model |
| 99 | + |
| 100 | + |
| 101 | +def _ipex_post_quant_process(model, example_inputs, inplace=False): |
| 102 | + """Convert to a jit model. |
| 103 | +
|
| 104 | + Args: |
| 105 | + model: a prepared model. |
| 106 | + example_inputs: used to trace torch model. |
| 107 | + inplace: whether to carry out model transformations in-place. |
| 108 | +
|
| 109 | + Returns: |
| 110 | + A converted jit model. |
| 111 | + """ |
| 112 | + model = ipex.quantization.convert(model, inplace=inplace) |
| 113 | + with torch.no_grad(): |
| 114 | + try: |
| 115 | + if isinstance(example_inputs, dict): |
| 116 | + model = torch.jit.trace(model, example_kwarg_inputs=example_inputs) |
| 117 | + else: |
| 118 | + model = torch.jit.trace(model, example_inputs) |
| 119 | + model = torch.jit.freeze(model.eval()) |
| 120 | + except: |
| 121 | + if isinstance(example_inputs, dict): |
| 122 | + model = torch.jit.trace(model, example_kwarg_inputs=example_inputs, strict=False, check_trace=False) |
| 123 | + else: |
| 124 | + model = torch.jit.trace(model, example_inputs, strict=False) |
| 125 | + model = torch.jit.freeze(model.eval()) |
| 126 | + # After freezing, run 1 time to warm up the profiling graph executor to insert prim::profile |
| 127 | + # At the 2nd run, the llga pass will be triggered and the model is turned into |
| 128 | + # an int8 model: prim::profile will be removed and will have LlgaFusionGroup in the graph |
| 129 | + simple_inference(model, example_inputs, iterations=2) |
| 130 | + return model |
0 commit comments