Directory structure: └── onnx-neural-compressor/ ├── docs/ │ ├── autotune.md │ ├── calibration.md │ ├── design.md │ ├── installation_guide.md │ ├── quantization.md │ ├── quantization_layer_wise.md │ ├── quantization_weight_only.md │ └── smooth_quant.md ├── examples/ │ ├── image_recognition/ │ │ └── resnet50/ │ │ └── quantization/ │ │ └── ptq_static/ │ │ ├── main.py │ │ ├── prepare_model.py │ │ ├── run_benchmark.sh │ │ └── run_quant.sh │ └── .config/ │ └── model_params_onnxrt.json └── onnx_neural_compressor/ ├── __init__.py ├── constants.py ├── data_reader.py ├── logger.py ├── onnx_model.py ├── utility.py ├── version.py ├── algorithms/ │ ├── __init__.py │ ├── utility.py │ ├── layer_wise/ │ │ ├── __init__.py │ │ └── core.py │ ├── post_training_quant/ │ │ ├── __init__.py │ │ ├── calibrate.py │ │ ├── calibrator.py │ │ └── operators/ │ │ ├── __init__.py │ │ ├── activation.py │ │ ├── argmax.py │ │ ├── attention.py │ │ ├── base_op.py │ │ ├── binary_op.py │ │ ├── concat.py │ │ ├── conv.py │ │ ├── direct_q8.py │ │ ├── embed_layernorm.py │ │ ├── gather.py │ │ ├── gavgpool.py │ │ ├── gemm.py │ │ ├── lstm.py │ │ ├── matmul.py │ │ ├── maxpool.py │ │ ├── pad.py │ │ ├── pooling.py │ │ ├── reduce.py │ │ ├── resize.py │ │ ├── split.py │ │ └── unary_op.py │ ├── smoother/ │ │ ├── __init__.py │ │ ├── calibrator.py │ │ └── core.py │ └── weight_only/ │ ├── __init__.py │ ├── awq.py │ ├── gptq.py │ └── rtn.py └── quantization/ ├── __init__.py ├── algorithm_entry.py ├── matmul_4bits_quantizer.py ├── matmul_nbits_quantizer.py ├── quant_utils.py ├── quantize.py └── tuning.py ================================================ FILE: docs/autotune.md ================================================ AutoTune ======================================== 1. [Overview](#overview) 2. [How it Works](#how-it-works) 3. [Working with Autotune](#working-with-autotune) 4. [Get Started](#get-started) ## Overview Neural Compressor aims to help users quickly deploy low-precision models by leveraging popular compression techniques, such as post-training quantization and weight-only quantization algorithms. Despite having a variety of these algorithms, finding the appropriate configuration for a model can be difficult and time-consuming. To address this, we built the `autotune` module which identifies the best algorithm configuration for models to achieve optimal performance under the certain accuracy criteria. This module allows users to easily use predefined tuning recipes and customize the tuning space as needed. ## How it Works The autotune module constructs the tuning space according to the pre-defined tuning set or users' tuning set. It iterates the tuning space and applies the configuration on given float model then records and compares its evaluation result with the baseline. The tuning process stops when meeting the exit policy. The workflow is as below: Workflow ## Working with Autotune The `autotune` API can be used across all algorithms supported by Neural Compressor. It accepts three primary arguments: `model_input`, `tune_config`, and `eval_fn`. The `TuningConfig` class defines the tuning process, including the tuning space, order, and exit policy. - Define the tuning space User can define the tuning space by setting `config_set` with an algorithm configuration or a set of configurations. ```python # Use the default tuning space config_set = config.get_woq_tuning_config() # Customize the tuning space with one algorithm configurations config_set = config.RTNConfig(weight_sym=False, weight_group_size=[32, 64]) # Customize the tuning space with two algorithm configurations config_set = [ config.RTNConfig(weight_sym=False, weight_group_size=32), config.GPTQConfig(weight_group_size=128, weight_sym=False), ] ``` - Define the tuning order The tuning order determines how the process traverses the tuning space and samples configurations. Users can customize it by configuring the `sampler`. Currently, we provide the [`default_sampler`](https://github.com/onnx/neural-compressor/blob/main/onnx_neural_compressor/quantization/tuning.py#L210), which samples configurations sequentially, always in the same order. - Define the exit policy The exit policy includes two components: accuracy goal (`tolerable_loss`) and the allowed number of trials (`max_trials`). The tuning process will stop when either condition is met. ## Get Started The example below demonstrates how to autotune a ONNX model on four `RTNConfig` configurations. ```python from onnx_neural_compressor.quantization import config, tuning def eval_fn(model) -> float: return ... tune_config = tuning.TuningConfig( config_set=config.RTNConfig( weight_sym=[False, True], weight_group_size=[32, 128] ), tolerable_loss=0.2, max_trials=10, ) q_model = tuning.autotune(model, tune_config=tune_config, eval_fn=eval_fn) ``` ================================================ FILE: docs/calibration.md ================================================ # Calibration Algorithms in Quantization 1. [Introduction](#introduction) 2. [Calibration Algorithms](#calibration-algorithms) 3. [Support Matrix](#support-matrix) ## Introduction Quantization proves beneficial in terms of reducing the memory and computational requirements of the model. Uniform quantization transforms the input value $x ∈ [β, α]$ to lie within $[−2^{b−1}, 2^{b−1} − 1]$, where $[β, α]$ is the range of real values chosen for quantization and $b$ is the bit-width of the signed integer representation. Calibration is the process of determining the $α$ and $β$ for model weights and activations. Refer to this [link](.quantization.md#quantization-fundamentals) for more quantization fundamentals ## Calibration Algorithms Currently, Neural Compressor supports three popular calibration algorithms: - MinMax: This method gets the maximum and minimum of input values as $α$ and $β$ [^1]. It preserves the entire range and is the simplest approach. - Entropy: This method minimizes the KL divergence to reduce the information loss between full-precision and quantized data [^2]. Its primary focus is on preserving essential information. - Percentile: This method only considers a specific percentage of values for calculating the range, ignoring the remainder which may contain outliers [^3]. It enhances resolution by excluding extreme values but still retaining noteworthy data. > `kl` is used to represent the Entropy calibration algorithm in Neural Compressor. ## Reference [^1]: Vanhoucke, Vincent, Andrew Senior, and Mark Z. Mao. "Improving the speed of neural networks on CPUs." (2011). [^2]: Szymon Migacz. "Nvidia 8-bit inference width tensorrt." (2017). [^3]: McKinstry, Jeffrey L., et al. "Discovering low-precision networks close to full-precision networks for efficient embedded inference." arXiv preprint arXiv:1809.04191 (2018). ================================================ FILE: docs/design.md ================================================ Design ===== Neural Compressor features an architecture and workflow that aids in increasing performance and faster deployments across infrastructures. ## Architecture Architecture ## Workflow Workflow ================================================ FILE: docs/installation_guide.md ================================================ # Installation 1. [Installation](#installation) 1.1. [Prerequisites](#prerequisites) 1.2. [Install from Binary](#install-from-binary) 1.3. [Install from Source](#install-from-source) 2. [System Requirements](#system-requirements) 2.1. [Validated Hardware Environment](#validated-hardware-environment) 2.2. [Validated Software Environment](#validated-software-environment) ## Installation ### Prerequisites You can install Neural Compressor from binary or source. The following prerequisites and requirements must be satisfied for a successful installation: - Python version: 3.8 or 3.9 or 3.10 or 3.11 ### Install from Binary ```Shell # install stable basic version from pypi pip install onnx-neural-compressor ``` ### Install from Source ```Shell git clone https://github.com/onnx/neural-compressor.git cd neural-compressor pip install -r requirements.txt pip install . ``` ## System Requirements ### Validated Hardware Environment #### Neural Compressor supports CPUs based on [Intel 64 architecture or compatible processors](https://en.wikipedia.org/wiki/X86-64): * Intel Xeon Scalable processor (formerly Skylake, Cascade Lake, Cooper Lake, Ice Lake, and Sapphire Rapids) * Intel Xeon CPU Max Series (formerly Sapphire Rapids HBM) * Intel Core Ultra Processors (Meteor Lake, Lunar Lake) ### Validated Software Environment * OS version: CentOS 8.4, Ubuntu 22.04 * Python version: 3.10 * ONNX Runtime version: 1.18.1 ================================================ FILE: docs/quantization.md ================================================ Quantization =============== 1. [Quantization Introduction](#quantization-introduction) 2. [Quantization Fundamentals](#quantization-fundamentals) 3. [Get Started](#get-started) 3.1 [Post Training Quantization](#post-training-quantization) 3.2 [Specify Quantization Rules](#specify-quantization-rules) 3.3 [Specify Quantization Backend and Device](#specify-quantization-backend-and-device) 4. [Examples](#examples) ## Quantization Introduction Quantization is a very popular deep learning model optimization technique invented for improving the speed of inference. It minimizes the number of bits required by converting a set of real-valued numbers into the lower bit data representation, such as int8 and int4, mainly on inference phase with minimal to no loss in accuracy. This way reduces the memory requirement, cache miss rate, and computational cost of using neural networks and finally achieve the goal of higher inference performance. On Intel 3rd Gen Intel® Xeon® Scalable Processors, user could expect up to 4x theoretical performance speedup. We expect further performance improvement with [Intel® Advanced Matrix Extensions](https://www.intel.com/content/www/us/en/products/docs/accelerator-engines/advanced-matrix-extensions/overview.html) on 4th Gen Intel® Xeon® Scalable Processors. ## Quantization Fundamentals `Affine quantization` and `Scale quantization` are two common range mapping techniques used in tensor conversion between different data types. The math equation is like: $X_{int8} = round(Scale \times X_{fp32} + ZeroPoint)$. **Affine Quantization** This is so-called `Asymmetric quantization`, in which we map the min/max range in the float tensor to the integer range. Here int8 range is [-128, 127], uint8 range is [0, 255]. here: If INT8 is specified, $Scale = (|X_{max} - X_{min}|) / 127$ and $ZeroPoint = -128 - X_{min} / Scale$. or If UINT8 is specified, $Scale = (|X_{max} - X_{min}|) / 255$ and $ZeroPoint = - X_{min} / Scale$. **Scale Quantization** This is so-called `Symmetric quantization`, in which we use the maximum absolute value in the float tensor as float range and map to the corresponding integer range. The math equation is like: here: If INT8 is specified, $Scale = max(abs(X_{max}), abs(X_{min})) / 127$ and $ZeroPoint = 0$. or If UINT8 is specified, $Scale = max(abs(X_{max}), abs(X_{min})) / 255$ and $ZeroPoint = 128$. *NOTE* Sometimes the reduce_range feature, that's using 7 bit width (1 sign bit + 6 data bits) to represent int8 range, may be needed on some early Xeon platforms, it's because those platforms may have overflow issues due to fp16 intermediate calculation result when executing int8 dot product operation. After AVX512_VNNI instruction is introduced, this issue gets solved by supporting fp32 intermediate data. ### Quantization Support Matrix | Framework | Backend Library | Symmetric Quantization | Asymmetric Quantization | | :-------------- |:---------------:| ---------------:|---------------:| | ONNX Runtime | [MLAS](https://github.com/microsoft/onnxruntime/tree/master/onnxruntime/core/mlas) | Activation (int8/uint8), Weight (int8/uint8) | Activation (int8/uint8), Weight (int8/uint8) | > ***Note*** > > Activation (uint8) + Weight (int8) is recommended for performance on x86-64 machines with AVX2 and AVX512 extensions. #### Reference + MLAS: [MLAS Quantization](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/quantization/onnx_quantizer.py) ### Quantization Approaches Quantization has two different approaches which belong to optimization on inference: 1) post training dynamic quantization 2) post training static quantization #### Post Training Dynamic Quantization The weights of the neural network get quantized into 8 bits format from float32 format offline. The activations of the neural network is quantized as well with the min/max range collected during inference runtime. This approach is widely used in dynamic length neural networks, like NLP model. #### Post Training Static Quantization Compared with `post training dynamic quantization`, the min/max range in weights and activations are collected offline on a so-called `calibration` dataset. This dataset should be able to represent the data distribution of those unseen inference dataset. The `calibration` process runs on the original fp32 model and dumps out all the tensor distributions for `Scale` and `ZeroPoint` calculations. Usually preparing 100 samples are enough for calibration. This approach is major quantization approach people should try because it could provide the better performance comparing with `post training dynamic quantization`. ## Get Started The design philosophy of the quantization interface of Neural Compressor is easy-of-use. It requests user to provide `model_input`, `model_output` and `quant_config`. Those parameters would be used to quantize and save the model. `model_input` is the ONNX model location or the ONNX model object. `model_output` is the path to save ONNX model. `quant_config` is the configuration to do quantization. User could leverage Neural Compressor to directly generate a fully quantized model without accuracy validation. Currently, Neural Compressor supports `Post Training Static Quantization` and `Post Training Dynamic Quantization`. ### Post Training Quantization ``` python from onnx_neural_compressor.quantization import quantize, config from onnx_neural_compressor import data_reader class DataReader(data_reader.CalibrationDataReader): def get_next(self): ... def rewind(self): ... calibration_data_reader = DataReader() # only needed by StaticQuantConfig qconfig = config.StaticQuantConfig(calibration_data_reader) # or qconfig = DynamicQuantConfig() quantize(model, q_model_path, qconfig) ``` ### Specify Quantization Rules Neural Compressor support specify quantization rules by operator name. Users can use `set_local` API of configs to achieve the above purpose by below code: ```python op_config = config.StaticQuantConfig(per_channel=False) quant_config = config.StaticQuantConfig( per_channel=True, ) quant_config.set_local("/h.4/mlp/fc_out/MatMul", op_config) ``` ### Specify Quantization Backend and Device Neural-Compressor will quantized models with user-specified backend or detecting the hardware and software status automatically to decide which backend should be used. The automatically selected priority is: GPU/NPU > CPU.
Backend Backend Library Support Device(cpu as default)
CPUExecutionProvider MLAS cpu
TensorrtExecutionProvider TensorRT gpu
CUDAExecutionProvider CUDA gpu
DnnlExecutionProvider OneDNN cpu
DmlExecutionProvider* OneDNN npu


> ***Note*** > > DmlExecutionProvider support works as experimental, please expect exceptions. > > Known limitation: the batch size of onnx models has to be fixed to 1 for DmlExecutionProvider, no multi-batch and dynamic batch support yet. ## Examples User could refer to [examples](../../examples) on how to quantize a new model. ================================================ FILE: docs/quantization_layer_wise.md ================================================ Layer Wise Quantization (LWQ) ===== 1. [Introduction](#introduction) 2. [Supported Framework Model Matrix](#supported-framework-model-matrix) 3. [Examples](#examples) ## Introduction Large language models (LLMs) have shown exceptional performance across various tasks, meanwhile, the substantial parameter size poses significant challenges for deployment. Layer-wise quantization(LWQ) can greatly reduce the memory footprint of LLMs, usually 80-90% reduction, which means that users can quantize LLMs even on single node using GPU or CPU. We can quantize the model under memory-constrained devices, therefore making the huge-sized LLM quantization possible. *Figure 1: The process of layer-wise quantization for ONNX model. The graph of LLM is split into several parts, and each subgraph is quantized in turn.* ## Supported Framework Model Matrix
Types/Framework ONNX Runtime
W8A8 Post Training Static Quantization ✕
Weight-only Quantization RTN ✔
AWQ ✕
GPTQ ✔
## Examples ```python from onnx_neural_compressor.quantization import matmul_4bits_quantizer algo_config = matmul_4bits_quantizer.RTNWeightOnlyQuantConfig(layer_wise_quant=True) quant = matmul_4bits_quantizer.MatMul4BitsQuantizer( model, algo_config=algo_config, ) quant.process() qmodel = quant.model ``` ================================================ FILE: docs/quantization_weight_only.md ================================================ Weight Only Quantization (WOQ) ===== 1. [Introduction](#introduction) 2. [Supported Framework Model Matrix](#supported-framework-model-matrix) 3. [Examples](#examples) 4. [WOQ Algorithms Tuning](#woq-algorithms-tuning) ## Introduction As large language models (LLMs) become more prevalent, there is a growing need for new and improved quantization methods that can meet the computational demands of these modern architectures while maintaining the accuracy. Compared to normal quantization like W8A8, weight only quantization is probably a better trade-off to balance the performance and the accuracy, since we will see below that the bottleneck of deploying LLMs is the memory bandwidth and normally weight only quantization could lead to better accuracy. Model inference: Roughly speaking , two key steps are required to get the model's result. The first one is moving the model from the memory to the cache piece by piece, in which, memory bandwidth $B$ and parameter count $P$ are the key factors, theoretically the time cost is $P*4 /B$. The second one is computation, in which, the device's computation capacity $C$ measured in FLOPS and the forward FLOPs $F$ play the key roles, theoretically the cost is $F/C$. Text generation: The most famous application of LLMs is text generation, which predicts the next token/word based on the inputs/context. To generate a sequence of texts, we need to predict them one by one. In this scenario, $F\approx P$ if some operations like bmm are ignored and past key values have been saved. However, the $C/B$ of the modern device could be to **100X,** that makes the memory bandwidth as the bottleneck in this scenario. Besides, as mentioned in many papers[1][2], activation quantization is the main reason to cause the accuracy drop. So for text generation task, weight only quantization is a preferred option in most cases. Theoretically, round-to-nearest (RTN) is the most straightforward way to quantize weight using scale maps. However, when the number of bits is small (e.g. 3), the MSE loss is larger than expected. A group size is introduced to reduce elements using the same scale to improve accuracy. There are many excellent works for weight only quantization to improve its accuracy performance, such as AWQ[3], GPTQ[4]. Neural compressor integrates these popular algorithms in time to help customers leverage them and deploy them to their own tasks. ## Supported Framework Model Matrix | Algorithms/Framework | ONNX Runtime | |--------------|----------| | RTN | ✔ | | AWQ | ✔ | | GPTQ | ✔ | > **RTN:** A quantification method that we can think of very intuitively. It does not require additional datasets and is a very fast quantization method. Generally speaking, RTN will convert the weight into a uniformly distributed integer data type, but some algorithms, such as Qlora, propose a non-uniform NF4 data type and prove its theoretical optimality. > **GPTQ:** A new one-shot weight quantization method based on approximate second-order information, that is both highly-accurate and highly efficient[4]. The weights of each column are updated based on the fixed-scale pseudo-quantization error and the inverse of the Hessian matrix calculated from the activations. The updated columns sharing the same scale may generate a new max/min value, so the scale needs to be saved for restoration. > **AWQ:** Proved that protecting only 1% of salient weights can greatly reduce quantization error. the salient weight channels are selected by observing the distribution of activation and weight per channel. The salient weights are also quantized after multiplying a big scale factor before quantization for preserving. ## Examples ### **Quantization Capability** | Config | Capability | |---|---| | weight_dtype | ['int'] | | weight_bits | [1, ..., 8] | | weight_group_size | [-1, 1, ..., $C_{in}$] | | weight_sym | ['asym', 'sym'] | | algorithm | ['RTN', 'AWQ', 'GPTQ'] | Notes: *weight_group_size = -1* refers to **per output channel quantization**. Taking a MatMul operator (input channel = $C_{in}$, output channel = $C_{out}$) for instance, when *weight_group_size = -1*, quantization will calculate total $C_{out}$ quantization parameters. Otherwise, when *weight_group_size = gs* quantization parameters are calculate with every $gs$ elements along with the input channel, leading to total $C_{out} \times (C_{in} / gs)$ quantization parameters. **RTN arguments** | rtn_args | default value | comments | |----------|-------------|-------------------------------------------------------------------| | accuracy_level | 0 | Support 0 (unset), 1(fp32 compute type of jblas kernel), 2 (fp16 compute type of jblas kernel), 3 (bf16 compute type of jblas kernel), 4 (int8 compute type of jblas kernel) | | ratios | {} | Percentile of clip | | providers | ["CPUExecutionProvider"] | Execution providers to use | | layer_wise_quant | False | Whether to quantize model layer by layer to save memory footprint. | | quant_last_matmul | True | Whether to quantize the last matmul of the model | **AWQ arguments** | awq_args | default value | comments | |----------|-------------|-------------------------------------------------------------------| | accuracy_level | 0 | Support 0 (unset), 1(fp32 compute type of jblas kernel), 2 (fp16 compute type of jblas kernel), 3 (bf16 compute type of jblas kernel), 4 (int8 compute type of jblas kernel) | | enable_auto_scale | True | Whether to search for best scales based on activation distribution | | enable_mse_search | True | Whether to search for the best clip range from range [0.91, 1.0, 0.01] | | providers | ["CPUExecutionProvider"] | Execution providers to use | | quant_last_matmul | True | Whether to quantize the last matmul of the model | **GPTQ arguments** | gptq_args | default value | comments | |----------|-------------|-------------------------------------------------------------------| | accuracy_level | 0 | Support 0 (unset), 1(fp32 compute type of jblas kernel), 2 (fp16 compute type of jblas kernel), 3 (bf16 compute type of jblas kernel), 4 (int8 compute type of jblas kernel) | | percdamp | 0.01 | Percentage of Hessian's diagonal values' average, which will be added to Hessian's diagonal to increase numerical stability| | block_size | 128 | Execute GPTQ quantization per block, block shape = [$C_{out}$, block_size] | | actorder | False | Whether to sort Hessian's diagonal values to rearrange channel-wise quantization order| | mse | False | Whether get scale and zero point with mse error | | perchannel | True | Whether quantize weight per-channel | | providers | ["CPUExecutionProvider"] | Execution providers to use | | layer_wise_quant | False | Whether to quantize model layer by layer to save memory footprint. | | quant_last_matmul | True | Whether to quantize the last matmul of the model | **Note:** Neural compressor provides `Unsigned integer for asymmetric quantization` and `Signed integer for symmetric quantization`. Please follow the below section to compress the low bit data type for saving. ### **User Code Example** ```python from onnx_neural_compressor.quantization import matmul_4bits_quantizer algo_config = matmul_4bits_quantizer.GPTQWeightOnlyQuantConfig(calibration_data_reader=calibration_data_reader) quant = matmul_4bits_quantizer.MatMul4BitsQuantizer( model, block_size=32, is_symmetric=False, algo_config=algo_config, ) quant.process() q_model = quant.model ``` ## WOQ Algorithms Tuning To find the best algorithm, users can leverage the `autotune` feature to explore a set of configurations. It automatically searches for the optimal one with the best result. Users have the option to specify their own list of potential configurations or utilize the pre-defined configuration set. **Pre-defined configurations** | WOQ configurations | Comments | |------------------|-------| |RTN_G32ASYM| {"algorithm": "RTN", "group_size": 32, "scheme": "asym"}| |GPTQ_G32ASYM| {"algorithm": "GPTQ", "group_size": 32, "scheme": "asym"}| |GPTQ_G32ASYM_DISABLE_LAST_MATMUL| {"algorithm": "GPTQ", "group_size": 32, "scheme": "asym"}
& disable last MatMul| |GPTQ_G128ASYM| {"algorithm": "GPTQ", "group_size": 128, "scheme": "asym"}| |AWQ_G32ASYM| {"algorithm": "AWQ", "group_size": 32, "scheme": "asym"}| ### **User code example** ```python from onnx_neural_compressor.quantization import tuning, config tune_config = tuning.TuningConfig(config_set=config.get_woq_tuning_config()) best_model = tuning.autotune( model_input=model, tune_config=tune_config, eval_fn=eval_fn, calibration_data_reader=data_reader, ) ``` Refer to this [link](../../examples/onnxrt/nlp/huggingface_model/text_generation/llama/quantization/weight_only) for an example of WOQ algorithms tuning on ONNX Llama models. ## Reference [1]. Xiao, Guangxuan, et al. "Smoothquant: Accurate and efficient post-training quantization for large language models." arXiv preprint arXiv:2211.10438 (2022). [2]. Wei, Xiuying, et al. "Outlier suppression: Pushing the limit of low-bit transformer language models." arXiv preprint arXiv:2209.13325 (2022). [3]. Lin, Ji, et al. "AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration." arXiv preprint arXiv:2306.00978 (2023). [4]. Frantar, Elias, et al. "Gptq: Accurate post-training quantization for generative pre-trained transformers." arXiv preprint arXiv:2210.17323 (2022). ================================================ FILE: docs/smooth_quant.md ================================================ # Smooth Quant 1. [Introduction](#Introduction) 2. [Quantization Fundamentals](#Quantization-Fundamentals) 3. [SmoothQuant and Our Enhancement](#SmoothQuant-and-Our-Enhancement) 4. [Validated Models](#Validated-Models) 5. [Usage](#Usage) ## Introduction Quantization is a common compression operation to reduce memory and accelerate inference by converting the floating point matrix to an integer matrix. For large language models (LLMs) with gigantic parameters, the systematic outliers make quantification of activations difficult. [SmoothQuant](https://arxiv.org/abs/2211.10438), a training free post-training quantization (PTQ) solution, offline migrates this difficulty from activations to weights with a mathematically equivalent transformation. ## Quantization Fundamentals Quantization is a common compression operation to reduce memory and accelerate inference; therefore, the difficulty of LLM deployment can be alleviated. Quantization converts the floating point matrix to an integer matrix. The equation of quantization is as follows: $$ X_{int8} = round(X_{fp32}/S) + Z \tag{1} $$ where $X_{fp32}$ is the input matrix, $S$ is the scale factor, $Z$ is the integer zero point. ### Per-tensor & Per-channel There are several choices of sharing quantization parameters among tensor elements, also called quantization granularity. The coarsest level, per-tensor granularity, is that all elements in the tensor share the same quantization parameters. Finer granularity means sharing quantization parameters per row or per column for 2D matrices and per channel for 3D matrices. Similarly, the finest granularity is that each element has an individual set of quantization parameters. However, due to the model accuracy and computational consumption, per-tensor or per-channel are usually adopted. **In the following part, We will show that per-channel could bring lower quantization loss but has some limitations, that is why normally we use per-channel for weight quantization and per-tensor for activation/input quantization** #### Per-tensor example Suppose the weight tensor is： ```python W = np.array( [ [0.6839, 0.4741, 0.7451], [0.9301, 0.1742, 0.6835], ] ) ``` According to the formula (1), we need scale $S$ and zero point $Z$ to calculate the integer matrix. $$ S = \frac{X_{max} - X{min}}{2^b -1} \tag{2} $$ $$ Z = -round(X_{min/}/S) \tag{3} $$ The per-tensor quantization function is: ```python def quantize(x, num_bits=8): q_min, q_max = 0, 2.0**num_bits - 1.0 scale = (np.max(x) - np.min(x)) / (2**num_bits - 1) scale = np.clip(scale, 1e-5, None) zp = (0 - (np.min(x)) / scale).round() q_x = x / scale + zp q_x = np.clip(q_x, q_min, q_max).round() print(f"scale = {scale}, zp = {zp}") return q_x, scale, zp ``` Then we can get the quantized $W_{q}$ ```bash >>> W_q, scale, zp = quantize(W) scale = 0.0029643137254901962, zp = -59.0 >>> W_q array([[172., 101., 192.], [255., 0., 172.]]) ``` With the value of scale and zp, we can dequantize the tensor. ```python def dequantize(q_x, scale, zp): return scale * (q_x - zp) ``` ```bash >>> W_dq = dequantize(W_q, 0.001, -50) >>> W_dq array([[0.222, 0.151, 0.242], [0.305, 0.05 , 0.222]]) >>> loss = ((W_dq - W)**2).mean() >>> loss 0.19833545500000002 >>> W_dq = dequantize(W_q, scale, zp) >>> W_dq array([[0.68475647, 0.4742902 , 0.74404275], [0.93079451, 0.17489451, 0.68475647]]) >>> loss = ((W_dq - W)**2).mean() >>> loss 7.384850698449426e-07 ``` The difference between $W$ and $W_{dq}$ shows that quantization affects precision and appropriate values of scale and zero point will reduce the loss of precision. #### Per-channel example Similarly, the example of per-channel quantization is as follows: ```python def quantize_per_channel(x, num_bits=8): q_min, q_max = 0, 2.0**num_bits - 1.0 x_tmp = np.reshape(x, (x.shape[0], -1)) scales = np.max(x_tmp, axis=-1, keepdims=True) / (2**num_bits - 1) zp = (0 - np.min(x_tmp, axis=-1, keepdims=True) / scales).round() q_x = x_tmp / scales + zp q_x = np.clip(q_x, q_min, q_max).round() print(f"scales = {scales}, \n zp = {zp}") return q_x, scales, zp def dequantize_per_channel(q_x, scales, zp): print(q_x, scales, zp) print(scales * (q_x - zp)) return scales * (q_x - zp) ``` ```bash >>>W_q, scales, zp = quantize_per_channel(W) scales = [[0.00292196] [0.00364745]], zp = [[-162.] [ -48.]] >>>W_q array([[ 72., 0., 93.], [207., 0., 139.]]) >>>W_dq = dequantize_per_channel(W_q, scales, zp) >>>W_dq [[0.68373882 0.47335765 0.7451 ] [0.9301 0.17507765 0.68207333]] ``` And the loss is ```bash >>> loss = ((W_dq - W)**2).mean() >>> loss.item() 5.637846469306487e-07 ``` Through this example, we can see that per-channel quantization has finer granularity and has lower loss (loss 5.6378e-07 for per-channel quantization and 7.3849e-07 for per-tensor quantization). #### Matmul quantization example For a MatMul in most model, $Y=X \cdot W$, we can quantize both the weights and activations in order to reduce the storage and accelerate inference. Using per-tensor scale quantization to show the process. ```python def quantize_per_tensor_absmax(x, n_bits=8): scales = np.max(np.abs(x)) q_max = 2 ** (n_bits - 1) - 1 scales = np.clip(scales, 1e-5, None) / q_max q_x = x / scales q_x = np.clip(q_x, -q_max, q_max).round() return q_x, scales def dequantize(q_x, scale): return scale * q_x ``` Randomly initialize the $W$ and $Y$, then calculate the result of $Y=X \cdot W$ ```bash >>>W = np.random.randn(2, 3).astype(np.float32) >>>X = np.random.randn(3, 4).astype(np.float32) >>>W array([[-0.75903535, -1.7662522 , 1.0559074 ], [ 0.47551736, 0.33230257, 0.63447773]], dtype=float32) >>>X array([[-0.9628984 , 0.5076066 , -0.54988813, 1.2411681 ], [-1.6626304 , 0.2284153 , 0.4905207 , -0.11352996], [ 1.3270313 , -0.78117365, -0.3452512 , 0.826362 ]], dtype=float32) >>>Y = np.matmul(W, X) >>>Y array([[ 5.068721 , -1.6135774 , -0.81355196, 0.13099377], [-0.16839942, -0.17835854, -0.31753427, 1.076779 ]], dtype=float32) ``` Quantize weight and activation, matmul(quantize(X), quantize(Y)) ```bash >>>W_q, W_scale = quantize_per_tensor_absmax(W) >>>X_q, X_scale = quantize_per_tensor_absmax(X) >>>print(f'{W_q}\n{W_scale}') >>>print(f'{X_q}\n{X_scale}') [[ -55. -127. 76.] [ 34. 24. 46.]] 0.013907497323404147 [[ -74. 39. -42. 95.] [-127. 17. 37. -9.] [ 101. -60. -26. 63.]] 0.013091578258304145 >>>Y_q = np.matmul(W_q, X_q) >>>Y_q array([[27875., -8864., -4365., 706.], [ -918., -1026., -1736., 5912.]], dtype=float32) >>>Y_dq = dequantize(Y_q, W_scale * X_scale) >>>Y_dq array([[ 5.0752316 , -1.6138781 , -0.7947403 , 0.12854218], [-0.16714126, -0.18680494, -0.3160754 , 1.0764043 ]], dtype=float32) ``` #### Per-channel limitation Though per-channel quantization could bring lower quantization error, we could not apply it for activations due to the difficulty of the dequantization. We would prove it in the following image and the zero point of quantization would be ignored for simplicity. The image on the left presents a normal MatMul forward with 1x2 input $x$ and 2x2 weight $w$. The results $y$ could be easily obtained by simple mathematics. In the middle image, we apply per-tensor quantization for activations and per-channel quantization for weights; the results after quantization that are denoted by $y_1$ and $y_2$, could be easily dequantized to the float results $y_{fp1}$ and $y_{fp2}$ by per channel scale $1.0/s_1s_x$ and $1.0/s_2s_x$. However, after applying per-channel quantization for activation (right image), we could not dequantize the $y_1$ and $y_2$ to float results.
## SmoothQuant and Our Enhancement ### SmoothQuant In the previous subsection, we have explained why per-channel quantization could not be applied for activation, even though it could lead to lower quantization loss. However, the quantization error loss of activation plays an important role in the accuracy loss of model quantization[^2][^3][^4]. To reduce the quantization loss of activations, lots of methods have been proposed. In the following, we briefly introduce SPIQ[^2], Outlier Suppression[^3] and Smoothquant[^4]. All these three methods share a similar idea to migrate the difficulty from activation quantization to weight quantization but differ in how much difficulty to be transferred. So **the first question is how to migrate the difficulty from activation to weights?** The solution is straightforward, that is to convert the network to an output equivalent network that is presented in the image below and apply quantization to this equivalent network. The intuition is that each channel of activation could be scaled to make it more quantization-friendly, similar to a fake per-channel activation quantization.
Please note that this conversion will make the quantization of weights more difficult, because the scales attached to weights shown above are per-input-channel, while quantization of weights is per-output-channel or per-tensor. So **the second question is how much difficulty to be migrated**, that is how to choose the **conversion per-channel scale** $s_{x1}$ and $s_{x2}$ from the above image. Different works adopt different ways. *SPIQ* just adopts the quantization scale of activations as the conversion per-channel scale. *Outlier suppression* adopts the scale of the preceding layernorm as the conversion per-channel scale. *Smoothquant* introduces a hyperparameter $\alpha$ as a smooth factor to calculate the conversion per-channel scale and balance the quantization difficulty of activation and weight. $$ s_j = max(|X_j|)^\alpha/max(|W_j|)^{1-\alpha} \tag{4} $$ j is the index of the input channels.
For most of the models such as OPT and BLOOM, $\alpha = 0.5$ is a well-balanced value to split the difficulty of weight and activation quantization. A larger $\alpha$ value could be used on models with more significant activation outliers to migrate more quantization difficulty to weights. ### Our enhancement: #### Algorithm: Auto-tuning of $\alpha$. SmoothQuant method aims to split the quantization difficulty of weight and activation by using a fixed-value $\alpha$ for an entire model. However, as the distributions of activation outliers vary not only across different models but also across different layers within a model, we hereby propose a method to obtain operator-wise optimal $\alpha$ values with the ability to tune automatically. Our proposed method consists of 8 major steps: - Hook input minimum and maximum values of operators to be smoothed. - Find a list of operators on which smoothquant could be performed. - Generate a list of $\alpha$ values of a user-defined range. - Calculate smoothing factor using $\alpha$ value, adjust parameters accordingly and forward the adjusted model given an input sample. - Perform per-channel quantization_dequantization of weights and per-tensor quantization_dequantization of activations to predict output. - Calculate the loss with respect to FP32 output, iterate the previous two steps given each $\alpha$ value and save the loss per alpha. - Apply criterion on input operator and obtain the optimal alpha values of a single input sample. - Iterate the previous three steps over a number of input samples and save the optimal $\alpha$ values. Multiple criteria (e.g min, max and mean) are supported to determine the $\alpha$ value. Both alpha range and criterion could be configured in AutoAlphaArgs. In our experiments, an $\alpha$ range of [0.0, 1.0] with a step_size of 0.1 is found to be well-balanced one for the majority of models. #### Engineering *fully automated*: users only need to pass a model and dataloader. ```python from onnx_neural_compressor.algorithms import Smoother smoother = Smoother( model, calibration_data_reader, providers=["CPUExecutionProvider"], ) smoothed_model = smoother.transform(alpha=0.7) # alpha could 'auto' to enable auto-tuning ``` *support lots of fusing patterns*: when applying the conversion per-channel scales, a mul layer needs to be inserted, which will introduce some overhead. The official code fuses this op to the previous layernorm, while we support more operator types like MatMul, Conv. Currently we only handle the operator whose scale could be fused, we are trying to support other operators, please stay tuned. ## Usage There are two ways to apply smooth quantization: 1) using a fixed `alpha` for the entire model or 2) determining the `alpha` through auto-tuning. ### Using a fixed `alpha` To set a fixed alpha for the entire model, users can follow this example: ```python from onnx_neural_compressor.quantization import config qconfig = config.StaticQuantConfig( data_reader, extra_options={"SmoothQuant": True, "SmoothQuantAlpha": 0.5, "SmoothQuantFolding": True} ) ``` Supported parameters description: "SmoothQuantAlpha": a float value. Default is 0.5. "SmoothQuantFolding": whether to fold mul into the previous operator if possible, where mul is required to update the input distribution during smoothing. ### Determining the `alpha` through auto-tuning Users can search for the best `alpha` at two levels: 1) for the entire model, and 2) for each operator. #### Auto-tune the `alpha` for the entire model The tuning process looks for the optimal `alpha` value from a list of `alpha` values provided by the user. > Please note that, it may a considerable amount of time as the tuning process applies each `alpha` to the entire model and uses the evaluation result on the entire dataset as the metric to determine the best `alpha`. Here is an example: ```python from onnx_neural_compressor.quantization import tuning, config qconfig = tuning.TuningConfig(config_set=[config.SmoothQuantConfig(alpha=np.arange(0.1, 0.5, 0.05).tolist())]) best_model = tuning.autotune( model_input=model, tune_config=qconfig, eval_fn=eval_fn, calibration_data_reader=data_reader, ) ``` #### Auto-tune the `alpha` for each operator In this case, the tuning process searches the optimal `alpha` of each operator by evaluating the loss with respect to FP32 output on a few batches of data. Here is an example: ```python from onnx_neural_compressor.quantization import quantize, config qconfig = config.StaticQuantConfig( data_reader, extra_options={ "SmoothQuant": True, "SmoothQuantAlpha": "auto", "SmoothQuantCalibIter": 1, "AutoAlphaArgs": { "alpha_min": 0.3, # min value of auto-tuning alpha search space "alpha_max": 0.7, # max value of auto-tuning alpha search space "alpha_step": 0.05, # step_size of auto-tuning alpha search space "attn_method": "min", }, }, ) quantize(model, output_model_path, qconfig) ``` ## Reference [^1]: Jason, Wei, et al. "Emergent Abilities of Large Language Models". Published in Transactions on Machine Learning Research (2022). [^2]: Yvinec, Edouard, et al. "SPIQ: Data-Free Per-Channel Static Input Quantization." Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2023. [^3]: Wei, Xiuying, et al. "Outlier suppression: Pushing the limit of low-bit transformer language models." arXiv preprint arXiv:2209.13325 (2022). [^4]: Xiao, Guangxuan, et al. "Smoothquant: Accurate and efficient post-training quantization for large language models." arXiv preprint arXiv:2211.10438 (2022). ================================================ FILE: examples/image_recognition/resnet50/quantization/ptq_static/main.py ================================================ # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint:disable=redefined-outer-name,logging-format-interpolation import argparse import collections import logging import os import re import time import cv2 import numpy as np import onnx import onnxruntime as ort from PIL import Image from sklearn import metrics from onnx_neural_compressor import data_reader, quantization from onnx_neural_compressor.quantization import config, tuning logger = logging.getLogger(__name__) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.WARN ) def _topk_shape_validate(preds, labels): # preds shape can be Nxclass_num or class_num(N=1 by default) # it's more suitable for 'Accuracy' with preds shape Nx1(or 1) output from argmax if isinstance(preds, int): preds = [preds] preds = np.array(preds) elif isinstance(preds, np.ndarray): preds = np.array(preds) elif isinstance(preds, list): preds = np.array(preds) preds = preds.reshape((-1, preds.shape[-1])) # consider labels just int value 1x1 if isinstance(labels, int): labels = [labels] labels = np.array(labels) elif isinstance(labels, tuple): labels = np.array([labels]) labels = labels.reshape((labels.shape[-1], -1)) elif isinstance(labels, list): if isinstance(labels[0], int): labels = np.array(labels) labels = labels.reshape((labels.shape[0], 1)) elif isinstance(labels[0], tuple): labels = np.array(labels) labels = labels.reshape((labels.shape[-1], -1)) else: labels = np.array(labels) # labels most have 2 axis, 2 cases: N(or Nx1 sparse) or Nxclass_num(one-hot) # only support 2 dimension one-shot labels # or 1 dimension one-hot class_num will confuse with N if len(preds.shape) == 1: N = 1 class_num = preds.shape[0] preds = preds.reshape([-1, class_num]) elif len(preds.shape) >= 2: N = preds.shape[0] preds = preds.reshape([N, -1]) class_num = preds.shape[1] label_N = labels.shape[0] assert label_N == N, "labels batch size should same with preds" labels = labels.reshape([N, -1]) # one-hot labels will have 2 dimension not equal 1 if labels.shape[1] != 1: labels = labels.argsort()[..., -1:] return preds, labels class TopK: def __init__(self, k=1): self.k = k self.num_correct = 0 self.num_sample = 0 def update(self, preds, labels, sample_weight=None): preds, labels = _topk_shape_validate(preds, labels) preds = preds.argsort()[..., -self.k :] if self.k == 1: correct = metrics.accuracy_score(preds, labels, normalize=False) self.num_correct += correct else: for p, l in zip(preds, labels): # get top-k labels with np.argpartition # p = np.argpartition(p, -self.k)[-self.k:] l = l.astype("int32") if l in p: self.num_correct += 1 self.num_sample += len(labels) def reset(self): self.num_correct = 0 self.num_sample = 0 def result(self): if self.num_sample == 0: logger.warning("Sample num during evaluation is 0.") return 0 return self.num_correct / self.num_sample class DataReader(data_reader.CalibrationDataReader): def __init__(self, model_path, dataset_location, image_list, batch_size=1, calibration_sampling_size=-1): self.batch_size = batch_size self.image_list = [] self.label_list = [] src_lst = [] label_lst = [] num = 0 with open(image_list, "r") as f: for s in f: image_name, label = re.split(r"\s+", s.strip()) src = os.path.join(dataset_location, image_name) if not os.path.exists(src): continue src_lst.append(src) label_lst.append(int(label)) if len(src_lst) == batch_size: self.image_list.append(src_lst) self.label_list.append(label_lst) num += batch_size if calibration_sampling_size > 0 and num >= calibration_sampling_size: break src_lst = [] label_lst = [] if len(src_lst) > 0: self.image_list.append(src_lst) self.label_list.append(label_lst) model = onnx.load(model_path, load_external_data=False) self.inputs_names = [input.name for input in model.graph.input] self.iter_next = iter(self.image_list) def _preprpcess(self, src): with Image.open(src) as image: image = np.array(image.convert("RGB")).astype(np.float32) image = image / 255.0 image = cv2.resize(image, (256, 256), interpolation=cv2.INTER_LINEAR) h, w = image.shape[0], image.shape[1] y0 = (h - 224) // 2 x0 = (w - 224) // 2 image = image[y0 : y0 + 224, x0 : x0 + 224, :] image = (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] image = image.transpose((2, 0, 1)) return image.astype("float32") def get_next(self): lst = next(self.iter_next, None) if lst is not None: return {self.inputs_names[0]: np.stack([self._preprpcess(src) for src in lst])} else: return None def rewind(self): self.iter_next = iter(self.image_list) def eval_func(model, dataloader, metric): metric.reset() sess = ort.InferenceSession(model, providers=ort.get_available_providers()) labels = dataloader.label_list for idx, batch in enumerate(dataloader): output = sess.run(None, batch) metric.update(output, labels[idx]) return metric.result() if __name__ == "__main__": logger.info("Evaluating ONNXRuntime full precision accuracy and performance:") parser = argparse.ArgumentParser( description="Resnet50 fine-tune examples for image classification tasks.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument("--model_path", type=str, help="Pre-trained model on onnx file") parser.add_argument("--dataset_location", type=str, help="Imagenet data path") parser.add_argument("--label_path", type=str, help="Imagenet label path") parser.add_argument("--benchmark", action="store_true", default=False) parser.add_argument("--tune", action="store_true", default=False, help="whether quantize the model") parser.add_argument("--output_model", type=str, help="output model path") parser.add_argument("--mode", type=str, help="benchmark mode of performance or accuracy") parser.add_argument( "--intra_op_num_threads", type=int, default=4, help="intra_op_num_threads for performance benchmark" ) parser.add_argument( "--quant_format", type=str, default="QOperator", choices=["QDQ", "QOperator"], help="quantization format" ) parser.add_argument( "--batch_size", default=1, type=int, ) args = parser.parse_args() top1 = TopK() dataloader = DataReader(args.model_path, args.dataset_location, args.label_path, args.batch_size) def eval(onnx_model): dataloader.rewind() return eval_func(onnx_model, dataloader, top1) if args.benchmark: if args.mode == "performance": total_time = 0.0 num_iter = 100 num_warmup = 10 sess_options = ort.SessionOptions() sess_options.intra_op_num_threads = args.intra_op_num_threads session = ort.InferenceSession(args.model_path, sess_options, providers=ort.get_available_providers()) ort_inputs = {} len_inputs = len(session.get_inputs()) inputs_names = [session.get_inputs()[i].name for i in range(len_inputs)] for idx, batch in enumerate(dataloader): if idx + 1 > num_iter: break tic = time.time() predictions = session.run(None, batch) toc = time.time() if idx >= num_warmup: total_time += toc - tic print("\n", "-" * 10, "Summary:", "-" * 10) print(args) throughput = (num_iter - num_warmup) / total_time print("Throughput: {} samples/s".format(throughput)) elif args.mode == "accuracy": acc_result = eval_func(args.model_path, dataloader, top1) print("Batch size = %d" % dataloader.batch_size) print("Accuracy: %.5f" % acc_result) if args.tune: calibration_data_reader = DataReader( args.model_path, args.dataset_location, args.label_path, args.batch_size, calibration_sampling_size=100 ) custom_tune_config = tuning.TuningConfig( config_set=config.StaticQuantConfig.get_config_set_for_tuning( quant_format=( quantization.QuantFormat.QOperator if args.quant_format == "QOperator" else quantization.QuantFormat.QDQ ), ) ) best_model = tuning.autotune( model_input=args.model_path, tune_config=custom_tune_config, eval_fn=eval, calibration_data_reader=calibration_data_reader, ) onnx.save(best_model, args.output_model) ================================================ FILE: examples/image_recognition/resnet50/quantization/ptq_static/prepare_model.py ================================================ import argparse import os import sys import urllib.request MODEL_URL = "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-v1-12.onnx" MAX_TIMES_RETRY_DOWNLOAD = 5 def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--input_model", type=str, required=False, default="resnet50-v1-12.onnx") parser.add_argument("--output_model", type=str, required=True) return parser.parse_args() def progressbar(cur, total=100): percent = "{:.2%}".format(cur / total) sys.stdout.write("\r[%-100s] %s" % ("#" * int(cur), percent)) sys.stdout.flush() def schedule(blocknum, blocksize, totalsize): if totalsize == 0: percent = 0 else: percent = min(1.0, blocknum * blocksize / totalsize) * 100 progressbar(percent) def download_model(url, model_name, retry_times=5): if os.path.isfile(model_name): print(f"{model_name} exists, skip download") return True print("download model...") retries = 0 while retries < retry_times: try: urllib.request.urlretrieve(url, model_name, schedule) break except KeyboardInterrupt: return False except: retries += 1 print(f"Download failed{', Retry downloading...' if retries < retry_times else '!'}") return retries < retry_times def prepare_model(input_model, output_model): # Download model from [ONNX Model Zoo](https://github.com/onnx/models) download_model(MODEL_URL, output_model, MAX_TIMES_RETRY_DOWNLOAD) if __name__ == "__main__": args = parse_arguments() prepare_model(args.input_model, args.output_model) ================================================ FILE: examples/image_recognition/resnet50/quantization/ptq_static/run_benchmark.sh ================================================ #!/bin/bash set -x function main { init_params "$@" run_benchmark } # init params function init_params { for var in "$@" do case $var in --input_model=*) input_model=$(echo "$var" |cut -f2 -d=) ;; --dataset_location=*) dataset_location=$(echo "$var" |cut -f2 -d=) ;; --label_path=*) label_path=$(echo "$var" |cut -f2 -d=) ;; --mode=*) mode=$(echo "$var" |cut -f2 -d=) ;; --intra_op_num_threads=*) intra_op_num_threads=$(echo "$var" |cut -f2 -d=) ;; esac done } # run_benchmark function run_benchmark { python main.py \ --model_path "${input_model}" \ --dataset_location "${dataset_location}" \ --label_path "${label_path-${dataset_location}/../val.txt}" \ --mode "${mode}" \ --batch_size 1 \ --intra_op_num_threads "${intra_op_num_threads-4}" \ --benchmark } main "$@" ================================================ FILE: examples/image_recognition/resnet50/quantization/ptq_static/run_quant.sh ================================================ #!/bin/bash set -x function main { init_params "$@" run_tuning } # init params function init_params { for var in "$@" do case $var in --input_model=*) input_model=$(echo "$var" |cut -f2 -d=) ;; --output_model=*) output_model=$(echo "$var" |cut -f2 -d=) ;; --dataset_location=*) dataset_location=$(echo "$var" |cut -f2 -d=) ;; --label_path=*) label_path=$(echo "$var" |cut -f2 -d=) ;; --quant_format=*) quant_format=$(echo "$var" |cut -f2 -d=) ;; esac done } # run_tuning function run_tuning { python main.py \ --model_path "${input_model}" \ --dataset_location "${dataset_location}" \ --label_path "${label_path-${dataset_location}/../val.txt}" \ --output_model "${output_model}" \ --quant_format "${quant_format-QOperator}" \ --tune } main "$@" ================================================ FILE: examples/.config/model_params_onnxrt.json ================================================ { "onnxrt": { "llama-2-7b-rtn": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf", "main_script": "main.py", "batch_size": 1, "algorithm": "RTN" }, "llama-2-7b-rtn-with-past": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "RTN" }, "llama-2-7b-rtn-with-past-qdq": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past-opset-21", "main_script": "main.py", "batch_size": 1, "algorithm": "RTN" }, "llama-2-7b-awq": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf", "main_script": "main.py", "batch_size": 1, "algorithm": "AWQ" }, "llama-2-7b-awq-with-past": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "AWQ" }, "llama-2-7b-awq-with-past-qdq": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past-opset-21", "main_script": "main.py", "batch_size": 1, "algorithm": "AWQ" }, "llama-2-7b-gptq": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf", "main_script": "main.py", "batch_size": 1, "algorithm": "GPTQ" }, "llama-2-7b-gptq-with-past": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "GPTQ" }, "llama-2-7b-gptq-with-past-qdq": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past-opset-21", "main_script": "main.py", "batch_size": 1, "algorithm": "GPTQ" }, "llama-2-7b-woq_tune": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf", "main_script": "main.py", "batch_size": 1, "algorithm": "WOQ_TUNE" }, "llama-2-7b-woq_tune-with-past": { "model_name": "meta-llama/Llama-2-7b-hf", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Llama-2-7b-hf-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "WOQ_TUNE" }, "llama-3-8b-gptq-with-past": { "model_name": "meta-llama/Meta-Llama-3-8B", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Meta-Llama-3-8B-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "GPTQ" }, "phi-3-mini-128k-instruct-rtn-with-past": { "model_name": "microsoft/Phi-3-mini-128k-instruct", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Phi-3-mini-128k-instruct-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "RTN" }, "qwen2-7b-instruct-rtn-with-past": { "model_name": "Qwen/Qwen2-7B-Instruct", "model_src_dir": "nlp/huggingface_model/text_generation/quantization/weight_only", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/Qwen2-7B-Instruct-with-past", "main_script": "main.py", "batch_size": 1, "algorithm": "RTN" }, "bert_base_MRPC": { "model_src_dir": "nlp/bert/quantization/ptq_static", "dataset_location": "/tf_dataset/pytorch/glue_data/MRPC", "input_model": "/tf_dataset2/models/onnx/bert_base_MRPC/bert.onnx", "main_script": "main.py", "batch_size": 8 }, "bert_base_MRPC_dynamic": { "model_src_dir": "nlp/bert/quantization/ptq_dynamic", "dataset_location": "/tf_dataset/pytorch/glue_data/MRPC", "input_model": "/tf_dataset2/models/onnx/bert_base_MRPC/bert.onnx", "main_script": "main.py", "batch_size": 8 }, "resnet50-v1-12_qdq": { "model_src_dir": "image_recognition/resnet50/quantization/ptq_static", "dataset_location": "/tf_dataset2/datasets/imagenet/ImagenetRaw/ILSVRC2012_img_val", "input_model": "/tf_dataset2/models/onnx/resnet50-v1-12/resnet50-v1-13.onnx", "main_script": "main.py", "batch_size": 1 }, "resnet50-v1-12": { "model_src_dir": "image_recognition/resnet50/quantization/ptq_static", "dataset_location": "/tf_dataset2/datasets/imagenet/ImagenetRaw/ILSVRC2012_img_val", "input_model": "/tf_dataset2/models/onnx/resnet50-v1-12/resnet50-v1-12.onnx", "main_script": "main.py", "batch_size": 1 }, "sd-v1-5-sq": { "model_src_dir": "nlp/huggingface_model/text_to_image/stable_diffusion_v1_5/quantization/ptq_static", "dataset_location": "", "input_model": "/tf_dataset2/models/onnx/sd_v1_5", "main_script": "main.py", "batch_size": 1 } } } ================================================ FILE: onnx_neural_compressor/__init__.py ================================================ # Copyright (c) 2021 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Neural Compressor: An open-source Python library supporting popular model compression techniques for ONNX models.""" ================================================ FILE: onnx_neural_compressor/constants.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Class for All constants.""" import datetime from packaging import version # constants for configs GLOBAL = "global" LOCAL = "local" DEFAULT_WHITE_LIST = "*" EMPTY_WHITE_LIST = None # config name BASE_CONFIG = "base_config" COMPOSABLE_CONFIG = "composable_config" RTN = "rtn" STATIC_QUANT = "static_quant" DYNAMIC_QUANT = "dynamic_quant" SMOOTH_QUANT = "smooth_quant" GPTQ = "gptq" AWQ = "awq" DEFAULT_WORKSPACE = "./nc_workspace/{}/".format(datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")) ONNXRT116_VERSION = version.Version("1.16.0") ONNXRT1161_VERSION = version.Version("1.16.1") ONNXRT119_VERSION = version.Version("1.19.0") PRIORITY_RTN = 60 PRIORITY_GPTQ = 70 PRIORITY_AWQ = 50 PRIORITY_SMOOTH_QUANT = 80 PRIORITY_STATIC_QUANT = 70 PRIORITY_DYNAMIC_QUANT = 60 MAXIMUM_PROTOBUF = 2147483648 WHITE_MODULE_LIST = ["MatMul", "Conv"] RTN_OP_LIST = ["MatMul"] AWQ_OP_LIST = ["MatMul"] GPTQ_OP_LIST = ["MatMul"] DYNAMIC_CPU_OP_LIST = ["FusedConv", "Conv", "EmbedLayerNormalization", "MatMul", "Gather", "Attention", "LSTM"] DYNAMIC_CUDA_OP_LIST = ["FusedConv", "Conv", "EmbedLayerNormalization", "MatMul", "Gather", "Attention", "LSTM"] DYNAMIC_DML_OP_LIST = [] DYNAMIC_DNNL_OP_LIST = ["FusedConv", "Conv", "EmbedLayerNormalization", "MatMul", "Gather", "Attention", "LSTM"] DYNAMIC_TRT_OP_LIST = [] STATIC_QDQ_CPU_OP_LIST = [ "FusedConv", "Conv", "Gather", "GatherElements", "GatherND", "Tile", "MatMul", "Gemm", "EmbedLayerNormalization", "Attention", "Relu", "Clip", "LeakyRelu", "Sigmoid", "MaxPool", "GlobalAveragePool", "Pad", "Split", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "Resize", "Abs", "Shrink", "Sign", "Flatten", "Expand", "Slice", "Mod", "ReduceMax", "ReduceMin", "CenterCropPad", ] STATIC_QDQ_CUDA_OP_LIST = [ "FusedConv", "Conv", "Gather", "MatMul", "Gemm", "EmbedLayerNormalization", "Attention", "Relu", "Clip", "LeakyRelu", "Sigmoid", "MaxPool", "GlobalAveragePool", "Pad", "Split", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "Resize", "Abs", "Shrink", "Sign", "Flatten", "Expand", "Slice", "Mod", "ReduceMax", "ReduceMin", ] STATIC_QDQ_DML_OP_LIST = [ "Conv", "MatMul", "Relu", "Clip", "MaxPool", ] STATIC_QDQ_DNNL_OP_LIST = [ "FusedConv", "Conv", "Gather", "MatMul", "Gemm", "EmbedLayerNormalization", "Attention", "Relu", "Clip", "LeakyRelu", "Sigmoid", "MaxPool", "GlobalAveragePool", "Pad", "Split", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "Resize", ] STATIC_QDQ_TRT_OP_LIST = [ "Conv", "MatMul", "Attention", "LeakyRelu", "Gather", "Sigmoid", "MaxPool", "EmbedLayerNormalization", "GlobalAveragePool", "Pad", "Split", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "Resize", "Gemm", "Add", ] STATIC_QOPERATOR_CPU_OP_LIST = [ "FusedConv", "Conv", "Gather", "GatherElements", "GatherND", "Tile", "MatMul", "Gemm", "EmbedLayerNormalization", "Attention", "Mul", "Relu", "Clip", "LeakyRelu", "Sigmoid", "MaxPool", "GlobalAveragePool", "Pad", "Split", "Add", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "ArgMax", "Resize", "Abs", "Shrink", "Sign", "Flatten", "Expand", "Slice", "Mod", "ReduceMax", "ReduceMin", "CenterCropPad", ] STATIC_QOPERATOR_CUDA_OP_LIST = [ "FusedConv", "Conv", "Gather", "MatMul", "Gemm", "EmbedLayerNormalization", "Attention", "Mul", "Relu", "Clip", "LeakyRelu", "Sigmoid", "MaxPool", "GlobalAveragePool", "Pad", "Split", "Add", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "ArgMax", "Resize", "Abs", "Shrink", "Sign", "Flatten", "Expand", "Slice", "Mod", "ReduceMax", "ReduceMin", ] STATIC_QOPERATOR_DML_OP_LIST = [ "Conv", "MatMul", "Mul", "Relu", "Clip", "MaxPool", "Add", ] STATIC_QOPERATOR_DNNL_OP_LIST = [ "FusedConv", "Conv", "Gather", "MatMul", "Gemm", "EmbedLayerNormalization", "Attention", "Mul", "Relu", "Clip", "LeakyRelu", "Sigmoid", "MaxPool", "GlobalAveragePool", "Pad", "Split", "Add", "Squeeze", "Reshape", "Concat", "AveragePool", "Unsqueeze", "Transpose", "ArgMax", "Resize", ] STATIC_QOPERATOR_TRT_OP_LIST = [] STATIC_QOPERATOR_OP_LIST_MAP = { "CPUExecutionProvider": STATIC_QOPERATOR_CPU_OP_LIST, "CUDAExecutionProvider": STATIC_QOPERATOR_CUDA_OP_LIST, "DmlExecutionProvider": STATIC_QOPERATOR_DML_OP_LIST, "DnnlExecutionProvider": STATIC_QOPERATOR_DNNL_OP_LIST, "TensorrtExecutionProvider": STATIC_QOPERATOR_TRT_OP_LIST, } STATIC_QDQ_OP_LIST_MAP = { "CPUExecutionProvider": STATIC_QDQ_CPU_OP_LIST, "CUDAExecutionProvider": STATIC_QDQ_CUDA_OP_LIST, "DmlExecutionProvider": STATIC_QDQ_DML_OP_LIST, "DnnlExecutionProvider": STATIC_QDQ_DNNL_OP_LIST, "TensorrtExecutionProvider": STATIC_QDQ_TRT_OP_LIST, } DYNAMIC_OP_LIST_MAP = { "CPUExecutionProvider": DYNAMIC_CPU_OP_LIST, "CUDAExecutionProvider": DYNAMIC_CUDA_OP_LIST, "DmlExecutionProvider": DYNAMIC_DML_OP_LIST, "DnnlExecutionProvider": DYNAMIC_DNNL_OP_LIST, "TensorrtExecutionProvider": DYNAMIC_TRT_OP_LIST, } ================================================ FILE: onnx_neural_compressor/data_reader.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import abc class CalibrationDataReader(metaclass=abc.ABCMeta): @classmethod def __subclasshook__(cls, subclass): return hasattr(subclass, "get_next") and callable(subclass.get_next) or NotImplemented @abc.abstractmethod def get_next(self) -> dict: """generate the input data dict for ONNXinferenceSession run""" raise NotImplementedError def __iter__(self): return self def __next__(self): result = self.get_next() if result is None: raise StopIteration return result @abc.abstractmethod def rewind(self): """Regenerate data.""" raise NotImplementedError ================================================ FILE: onnx_neural_compressor/logger.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging import os def _pretty_dict(value, indent=0): """Make the logger dict pretty.""" prefix = "\n" + " " * (indent + 4) if isinstance(value, dict): items = [prefix + repr(key) + ": " + _pretty_dict(value[key], indent + 4) for key in value] return "{%s}" % (",".join(items) + "\n" + " " * indent) elif isinstance(value, list): items = [prefix + _pretty_dict(item, indent + 4) for item in value] return "[%s]" % (",".join(items) + "\n" + " " * indent) elif isinstance(value, tuple): items = [prefix + _pretty_dict(item, indent + 4) for item in value] return "(%s)" % (",".join(items) + "\n" + " " * indent) else: return repr(value) LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper() _logger = logging.getLogger("onnx_neural_compressor") _logger.handlers.clear() _logger.setLevel(LOGLEVEL) formatter = logging.Formatter("%(asctime)s [%(levelname)s][%(filename)s:%(lineno)d] %(message)s", "%Y-%m-%d %H:%M:%S") streamHandler = logging.StreamHandler() streamHandler.setFormatter(formatter) _logger.addHandler(streamHandler) _logger.propagate = False def log(level, msg, *args, **kwargs): """Output log with the level as a parameter.""" kwargs.setdefault("stacklevel", 2) if isinstance(msg, dict): for _, line in enumerate(_pretty_dict(msg).split("\n")): _logger.log(level, line, *args, **kwargs) else: _logger.log(level, msg, *args, **kwargs) def debug(msg, *args, **kwargs): """Output log with the debug level.""" kwargs.setdefault("stacklevel", 2) if isinstance(msg, dict): for _, line in enumerate(_pretty_dict(msg).split("\n")): _logger.debug(line, *args, **kwargs) else: _logger.debug(msg, *args, **kwargs) def error(msg, *args, **kwargs): """Output log with the error level.""" kwargs.setdefault("stacklevel", 2) if isinstance(msg, dict): for _, line in enumerate(_pretty_dict(msg).split("\n")): _logger.error(line, *args, **kwargs) else: _logger.error(msg, *args, **kwargs) def fatal(msg, *args, **kwargs): """Output log with the fatal level.""" kwargs.setdefault("stacklevel", 2) if isinstance(msg, dict): for _, line in enumerate(_pretty_dict(msg).split("\n")): _logger.fatal(line, *args, **kwargs) else: _logger.fatal(msg, *args, **kwargs) def info(msg, *args, **kwargs): """Output log with the info level.""" kwargs.setdefault("stacklevel", 2) if isinstance(msg, dict): for _, line in enumerate(_pretty_dict(msg).split("\n")): _logger.info(line, *args, **kwargs) else: _logger.info(msg, *args, **kwargs) def warning(msg, *args, **kwargs): """Output log with the warning level (Alias of the method warn).""" kwargs.setdefault("stacklevel", 2) if isinstance(msg, dict): for _, line in enumerate(_pretty_dict(msg).split("\n")): _logger.warning(line, *args, **kwargs) else: _logger.warning(msg, *args, **kwargs) ================================================ FILE: onnx_neural_compressor/onnx_model.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Class for ONNX model.""" import collections import copy import os import pathlib import sys import onnx import transformers from onnx_neural_compressor import constants, logger, utility class ONNXModel: """Build ONNX model.""" def __init__(self, model, **kwargs): """Initialize an ONNX model. Args: model (str or ModelProto): path to onnx model or loaded ModelProto model object. """ self.model = model if not isinstance(model, str) else onnx.load(model, load_external_data=False) self._model_path = None if not isinstance(model, str) else model self.check_is_large_model() if self._is_large_model and self._model_path is None and not kwargs.get("ignore_warning", False): logger.warning("Model size > 2GB. Please use model path instead of onnx model object to quantize") if self._is_large_model and isinstance(model, str) and kwargs.get("load_external_data", True): onnx.external_data_helper.load_external_data_for_model(self.model, os.path.dirname(self._model_path)) self._config = None if isinstance(model, str) and os.path.exists(pathlib.Path(model).parent.joinpath("config.json").as_posix()): self._config = transformers.PretrainedConfig.from_pretrained(pathlib.Path(model).parent.as_posix()) self.node_name_counter = {} self._output_name_to_node = {} self._input_name_to_nodes = {} self._get_output_name_to_node(self.model.graph.node) self._get_input_name_to_nodes(self.model.graph.node) self._graph_info = {} self._get_graph_info() self._q_config = None def output_name_to_node(self): self._output_name_to_node = {} self._get_output_name_to_node(self.model.graph.node) return self._output_name_to_node def input_name_to_nodes(self): self._input_name_to_nodes = {} self._get_input_name_to_nodes(self.model.graph.node) return self._input_name_to_nodes def _get_input_name_to_nodes(self, nodes): """Get input names of nodes.""" for node in nodes: attrs = [ attr for attr in node.attribute if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS ] if len(attrs) > 0: for attr in attrs: self._get_input_name_to_nodes(attr.g.node) for input_name in node.input: if len(input_name.strip()) != 0: if input_name not in self._input_name_to_nodes: self._input_name_to_nodes[input_name] = [node] else: self._input_name_to_nodes[input_name].append(node) def _get_output_name_to_node(self, nodes): """Get output names of nodes.""" for node in nodes: attrs = [ attr for attr in node.attribute if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS ] if len(attrs) > 0: for attr in attrs: self._get_output_name_to_node(attr.g.node) for output_name in node.output: if len(output_name.strip()) != 0: self._output_name_to_node[output_name] = node @property def model_path(self): """Return model path.""" return self._model_path @model_path.setter def model_path(self, path): """Set model path.""" self._model_path = path def check_is_large_model(self): """Check model > 2GB.""" init_size = 0 for init in self.model.graph.initializer: # if initializer has external data location, return True if init.HasField("data_location") and init.data_location == onnx.TensorProto.EXTERNAL: self._is_large_model = True return # if raise error of initializer size > 2GB, return True try: init_bytes = init.SerializeToString() init_size += sys.getsizeof(init_bytes) except Exception as e: if "exceeds maximum protobuf size of 2GB" in str(e): self._is_large_model = True return else: # pragma: no cover raise e if init_size > constants.MAXIMUM_PROTOBUF: self._is_large_model = True return self._is_large_model = False @property def is_large_model(self): """Check the onnx model is over 2GB.""" return self._is_large_model @property def framework(self): """Return framework.""" return "onnxruntime" def add_initializer(self, tensor): """Add a initializer to model.""" if tensor.name not in [i.name for i in self._model.graph.initializer]: self._model.graph.initializer.append(tensor) def add_initializers(self, tensors): """Add initializers to model.""" for tensor in tensors: self.add_initializer(tensor) @property def q_config(self): """Return q_config.""" return self._q_config @q_config.setter def q_config(self, q_config): """Set q_config.""" self._q_config = q_config @property def hf_config(self): """Return huggingface config if model is Transformer-based.""" return self._config def input(self): """Return input of model.""" return [i.name for i in self.model.graph.input] def output(self): """Return output of model.""" return [i.name for i in self.model.graph.output] @property def model(self): """Return model itself.""" return self._model @model.setter def model(self, model): """Set model itself.""" self._model = model self._graph_info = {} self._get_graph_info() self._output_name_to_node = {} self._input_name_to_nodes = {} self._get_input_name_to_nodes(self._model.graph.node) self._get_output_name_to_node(self._model.graph.node) def nodes(self): """Return model nodes.""" return self._model.graph.node def initializer(self): """Return model initializer.""" return self._model.graph.initializer def graph(self): """Return model graph.""" return self._model.graph @property def ir_version(self): """Return model ir_version.""" return self._model.ir_version @property def opset_import(self): """Return model opset_import.""" return self._model.opset_import def update(self): """Update model info.""" self._graph_info = {} self._get_graph_info() self._output_name_to_node = self.output_name_to_node() self._input_name_to_nodes = self.input_name_to_nodes() @property def graph_info(self): """Return ORT Graph Info object holding information about backend graph.""" return self._graph_info def _get_graph_info(self): """Update graph info.""" for node in self.model.graph.node: self.graph_info.update({node.name: node.op_type}) def is_graph_output(self, name): """Check whether the tensor is the graph output.""" return name in self.output() def save(self, root): """Save ONNX model.""" if os.path.split(root)[0] != "" and not os.path.exists(os.path.split(root)[0]): os.mkdir(os.path.split(root)[0]) if self.is_large_model: # pragma: no cover onnx.external_data_helper.load_external_data_for_model(self.model, os.path.split(self._model_path)[0]) onnx.save_model( self.model, root, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(root) + "_data", size_threshold=1024, convert_attribute=False, ) else: onnx.save(self.model, root) self._model_path = root if self._config is not None and not os.path.exists(os.path.join(os.path.split(root)[0], "config.json")): model_type = "" if not hasattr(self._config, "model_type") else getattr(self._config, "model_type") setattr(self._config.__class__, "model_type", model_type) output_config_file = pathlib.Path(root).parent.joinpath("config.json").as_posix() self._config.to_json_file(output_config_file, use_diff=False) def remove_initializer(self, tensor): """Remove an initializer from model.""" if tensor in self._model.graph.initializer: self._model.graph.initializer.remove(tensor) def remove_initializers(self, init_to_remove): """Remove initializers from model.""" for initializer in init_to_remove: self.remove_initializer(initializer) def get_initializer(self, name): """ "Find the initializer with specified name.""" for initializer in self.model.graph.initializer: if initializer.name == name: return initializer return None def remove_node(self, node): """Remove a node from model.""" if node in self._model.graph.node: self._model.graph.node.remove(node) def remove_nodes(self, nodes_to_remove): """Remove nodes from model.""" for node in nodes_to_remove: self.remove_node(node) def add_node(self, node): """Add a node to model.""" self._model.graph.node.extend([node]) def add_nodes(self, nodes_to_add): """Add nodes to model.""" self._model.graph.node.extend(nodes_to_add) def get_children(self, node, input_name_to_nodes=None): """Get children nodes.""" if input_name_to_nodes is None: input_name_to_nodes = self._input_name_to_nodes children = [] for output in node.output: if output in input_name_to_nodes: for child in input_name_to_nodes[output]: children.append(child) return children def get_initializer_share_num(self, name): """Get the number of shares of initializer.""" num = 0 if self.get_initializer(name) is None: return num for node in self.nodes(): if name in node.input: num += 1 return num def get_node(self, name): """Get a node by name.""" for node in self.model.graph.node: if node.name == name: return node return None def get_parent(self, node, idx, output_name_to_node=None): if output_name_to_node is None: output_name_to_node = self._output_name_to_node if len(node.input) <= idx: return None input = node.input[idx] return output_name_to_node.get(input, None) def get_parents(self, node, output_name_to_node=None): if output_name_to_node is None: output_name_to_node = self._output_name_to_node parents = [] for input in node.input: if input in output_name_to_node: parents.append(output_name_to_node[input]) return parents def get_node_by_weight(self, weight_name): """Get a node by its weight name.""" if len(self._input_name_to_nodes) == 0: self._input_name_to_nodes = self.input_name_to_nodes() nodes = self._input_name_to_nodes[weight_name] if len(nodes) == 1: return nodes[0] elif len(nodes) == 0: raise ValueError("{} is not used by any node in this model.".format(weight_name)) else: raise NotImplementedError("Models with shared weights is not supported.") def set_initializer(self, tensor, array, raw=False): """Update initializer.""" old_tensor = self.get_initializer(tensor) self.remove_initializer(old_tensor) dims = old_tensor.dims data_type = old_tensor.data_type new_tensor = ( onnx.helper.make_tensor(tensor, data_type, dims, array.flatten().tolist()) if not raw else onnx.helper.make_tensor(tensor, data_type, dims, array.tostring(), raw=raw) ) self.add_initializer(new_tensor) def get_siblings(self, node): """Get siblings nodes.""" siblings = [] for parent in self.get_parents(node): for child in self.get_children(parent): if child.name != node.name: siblings.append(child) return siblings def get_scale_zero(self, tensor): """Help function to get scale and zero_point.""" if not tensor.endswith("_quantized"): logger.debug("Find {} in the quantized graph is not quantized.".format(tensor)) return None, None if len(self._input_name_to_nodes) == 0: self._input_name_to_nodes = self.input_name_to_nodes() if len(self._output_name_to_node) == 0: self._output_name_to_node = self.output_name_to_node() def _searcher(tensor_name): """Search scale and zero point tensor recursively.""" node = self._input_name_to_nodes[tensor_name][0] parent = self._output_name_to_node[tensor_name] if tensor_name in self._output_name_to_node else None direct_int8 = ["Reshape", "Transpose", "Squeeze", "Unsqueeze", "MaxPool", "Pad", "Split"] if parent is not None and parent.op_type in direct_int8: fp32_tensor_name = ( parent.input[0] .replace("_quantized", "") .replace("_QuantizeLinear", "") .replace("_QuantizeInput", "") ) elif node.op_type in ["Gather"]: # pragma: no cover fp32_tensor_name = ( node.output[0] .replace("_quantized", "") .replace("_QuantizeLinear", "") .replace("_QuantizeInput", "") ) else: fp32_tensor_name = ( tensor_name.replace("_quantized", "").replace("_QuantizeLinear", "").replace("_QuantizeInput", "") ) scale = fp32_tensor_name + "_scale" scale_tensor = self.get_initializer(scale) zo = fp32_tensor_name + "_zero_point" zo_tensor = self.get_initializer(zo) if scale_tensor is None or zo_tensor is None: if parent is not None: scale_tensor, zo_tensor = _searcher(parent.input[0]) return scale_tensor, zo_tensor node = self._input_name_to_nodes[tensor][0] # TODO check if scale_tensor and zero_point is needed # for bias of qlinearconv, scale and zero_point is not needed if (node.op_type == "QLinearConv" and tensor == node.input[-1]) or ( node.op_type == "QGemm" and tensor == node.input[-3] ): return None, None else: scale_tensor, zo_tensor = _searcher(tensor) assert scale_tensor, "missing scale for tensor {}".format(tensor) assert zo_tensor, "missing zero point for tensor {}".format(tensor) return scale_tensor, zo_tensor @staticmethod def replace_node_input(node, old_input_name, new_input_name): """Replace input of a node.""" assert isinstance(old_input_name, str) and isinstance(new_input_name, str) for j in range(len(node.input)): if node.input[j] == old_input_name: node.input[j] = new_input_name @staticmethod def replace_node_output(node, old_output_name, new_output_name): """Replace output of a node.""" assert isinstance(old_output_name, str) and isinstance(new_output_name, str) for j in range(len(node.output)): if node.output[j] == old_output_name: node.output[j] = new_output_name def replace_input_of_all_nodes(self, old_input_name, new_input_name, white_optype=[], black_optype=[]): """Replace inputs of all nodes.""" if len(white_optype) > 0: for node in self.model.graph.node: if node.op_type in white_optype: ONNXModel.replace_node_input(node, old_input_name, new_input_name) else: for node in self.model.graph.node: if node.op_type not in black_optype: ONNXModel.replace_node_input(node, old_input_name, new_input_name) def replace_output_of_all_nodes(self, old_output_name, new_output_name, white_optype=[], black_optype=[]): """Replace outputs of all nodes.""" if len(white_optype) > 0: for node in self.model.graph.node: if node.op_type in white_optype: ONNXModel.replace_node_output(node, old_output_name, new_output_name) else: for node in self.model.graph.node: if node.op_type not in black_optype: ONNXModel.replace_node_output(node, old_output_name, new_output_name) def remove_unused_nodes(self): """Remove unused nodes.""" unused_nodes = [] for node in self.model.graph.node: # remove constant if node.op_type == "Constant": tensor = node.attribute[0].t tensor.name = node.output[0] self.add_initializer(tensor) unused_nodes.append(node) # remove identity if node.op_type == "Identity": tensor = self.get_initializer(node.input[0]) if tensor is not None: new_tensor = copy.deepcopy(tensor) new_tensor.name = node.output[0] unused_nodes.append(node) self.add_initializer(new_tensor) self.remove_nodes(unused_nodes) if len(self._input_name_to_nodes) == 0: self._input_name_to_nodes = self.input_name_to_nodes() if len(self._output_name_to_node) == 0: self._output_name_to_node = self.output_name_to_node() unvalid_nodes = [ i for i in self.model.graph.node if all(out not in self._input_name_to_nodes and out not in self.output() for out in i.output) ] while len(unvalid_nodes) > 0: self.remove_nodes(unvalid_nodes) self._input_name_to_nodes = self.input_name_to_nodes() unvalid_nodes = [ i for i in self.model.graph.node if all([out not in self._input_name_to_nodes and out not in self.output() for out in i.output]) ] ununsed_weights = [] for w in self.model.graph.initializer: if w.name not in self._input_name_to_nodes and w.name not in self.output(): ununsed_weights.append(w) # Remove from graph.input for graph_input in self.graph().input: if graph_input.name == w.name: self.graph().input.remove(graph_input) self.remove_initializers(ununsed_weights) self.update() self.topological_sort() def topological_sort(self, enable_subgraph=False): """Topological sort the model.""" if not enable_subgraph: input_name_to_nodes = {} output_name_to_node = {} for node in self.model.graph.node: for input_name in node.input: if len(input_name.strip()) != 0: if input_name not in input_name_to_nodes: input_name_to_nodes[input_name] = [node] else: input_name_to_nodes[input_name].append(node) for output_name in node.output: if len(output_name.strip()) != 0: output_name_to_node[output_name] = node else: # pragma: no cover if len(self._input_name_to_nodes) == 0: self._input_name_to_nodes = self.input_name_to_nodes() if len(self._output_name_to_node) == 0: self._output_name_to_node = self.output_name_to_node() input_name_to_nodes = self._input_name_to_nodes output_name_to_node = self._output_name_to_node all_nodes = {} q = collections.deque() wait = collections.deque() for inp in self.model.graph.input: q.extend(input_name_to_nodes[inp.name]) for n in self.model.graph.node: if all([i not in output_name_to_node and i not in self.input() for i in n.input]): q.append(n) while q: n = q.popleft() if not all([output_name_to_node[i].name in all_nodes for i in n.input if i in output_name_to_node]): if n not in wait: wait.append(n) continue all_nodes[n.name] = n for out in n.output: if out in input_name_to_nodes: q.extend([i for i in input_name_to_nodes[out] if i.name not in all_nodes and i not in q]) if len(q) == 0 and len(wait) != 0: q = copy.deepcopy(wait) wait.clear() nodes = [i[1] for i in all_nodes.items()] assert len(list(set([n.name for n in nodes]))) == len(list(set([n.name for n in self.model.graph.node]))) self.model.graph.ClearField("node") self.model.graph.node.extend(nodes) def add_tensors_to_outputs(self, tensor_names): """Add the tensors to the model outputs to gets their values. Args: tensor_names: The names of tensors to be dumped. """ added_outputs = [] for tensor in tensor_names: if tensor not in self.output(): added_tensor = onnx.helper.ValueInfoProto() added_tensor.name = tensor added_outputs.append(added_tensor) self.model.graph.output.extend(added_outputs) # pylint: disable=no-member def remove_tensors_from_outputs(self, tensor_names): """Remove the tensors from the model outputs. Args: tensor_names: The names of tensors to be removed. """ removed_outputs = [] for tensor in tensor_names: if tensor in self.output(): removed_outputs.append(self.model.graph.output[self.output().index(tensor)]) for output in removed_outputs: self.model.graph.output.remove(output) def match_first_parent(self, node, parent_op_type, output_name_to_node_dict, exclude=[]): """Find parent node based on constraints on op_type. Args: node (str): current node name. parent_op_type (str): constraint of parent node op_type. output_name_to_node (dict): dictionary with output name as key, and node as value. exclude (list): list of nodes that are excluded (not allowed to match as parent). Returns: parent: The matched parent node. None if not found. index: The input index of matched parent node. None if not found. """ for i, input in enumerate(node.input): if input in output_name_to_node_dict: parent = output_name_to_node_dict[input] if parent.op_type == parent_op_type and parent not in exclude: return parent, i return None, None def match_parent( self, node, parent_op_type, input_index=None, output_name_to_node_dict=None, exclude=[], return_indice=None, ): """Find parent node based on constraints on op_type and index. Args: node (str): current node name. parent_op_type (str): constraint of parent node op_type. input_index (int or None): only check the parent given input index of current node. output_name_to_node (dict): dictionary with output name as key, and node as value. exclude (list): list of nodes that are excluded (not allowed to match as parent). return_indice (list): a list to append the input index when input_index is None. Returns: parent: The matched parent node. """ assert node is not None assert input_index is None or input_index >= 0 if output_name_to_node_dict is None: if len(self._output_name_to_node) == 0: self._output_name_to_node = self.output_name_to_node() output_name_to_node_dict = self._output_name_to_node if input_index is None: parent, index = self.match_first_parent(node, parent_op_type, output_name_to_node_dict, exclude) if return_indice is not None: return_indice.append(index) return parent if input_index >= len(node.input): return None parent = self.get_parent(node, input_index, output_name_to_node_dict) if parent is not None and parent.op_type == parent_op_type and parent not in exclude: return parent return None def match_parent_path( self, node, parent_op_types, parent_input_index, output_name_to_node_dict=None, return_indice=None, ): """Find a sequence of input edges based on constraints on parent op_type and index. Args: node (str): current node name. parent_op_types (str): constraint of parent node op_type of each input edge. parent_input_index (list): constraint of input index of each input edge. None means no constraint. output_name_to_node (dict): dictionary with output name as key, and node as value. return_indice (list): a list to append the input index when there is no constraint on input index of an edge. Returns: parents: a list of matched parent node. """ assert len(parent_input_index) == len(parent_op_types) if output_name_to_node_dict is None: if len(self._output_name_to_node) == 0: self._output_name_to_node = self.output_name_to_node() output_name_to_node_dict = self._output_name_to_node current_node = node matched_parents = [] for i, op_type in enumerate(parent_op_types): matched_parent = self.match_parent( current_node, op_type, parent_input_index[i], output_name_to_node_dict, exclude=[], return_indice=return_indice, ) if matched_parent is None: return None matched_parents.append(matched_parent) current_node = matched_parent return matched_parents def is_smoothquant_model(self): """Check the model is smooth quantized or not. Returns: bool: the model is smooth quantized or not. """ for init in self.model.graph.initializer: if "_smooth_scale" in init.name: return True return False # below functions are used for layer-wise def find_split_node_for_layer_wise_quantization(self): """Find split node for layer wise quantization.""" # find split nodes of decoder blocks # embed -> decoder.0 -(split_node)-> ... -(split_node)-> decoder.n -(split_node)-> norm -> head # after split: embed -> decoder.0, # decoder.1, # decoder.2, # ..., # decoder.n, # norm -> head start_nodes = [] for node in self.model.graph.node: start_node, qkv_nodes_list = None, None if node.op_type == "SkipLayerNormalization": start_node = node qkv_nodes_list = [ self.match_parent_path( start_node, ["MatMul", "Reshape", "Transpose", "Reshape", "MatMul"], [None, 0, 0, 0, 0], ), self.match_parent_path( start_node, ["Add", "MatMul", "Reshape", "Transpose", "MatMul"], [1, 1, 0, 0, 0], ), ] if node.op_type == "Add": start_node = node qkv_nodes_list = [ # match base attention structure self.match_parent_path( start_node, ["Add", "MatMul", "Reshape", "Transpose", "MatMul"], [0, None, 0, 0, 0], ), self.match_parent_path( start_node, ["Add", "MatMul", "Reshape", "Transpose", "MatMul"], [1, None, 0, 0, 0] ), # match gpt attention no past structure self.match_parent_path( start_node, ["Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"], [None, 0, 0, 0, 0, 0], output_name_to_node_dict=self._output_name_to_node, return_indice=[], ), # match bart attention structure self.match_parent_path( start_node, ["Add", "MatMul", "Reshape", "Transpose", "Reshape", "MatMul"], [0, None, 0, 0, 0, 0], ), self.match_parent_path( start_node, ["Add", "MatMul", "Reshape", "Transpose", "Reshape", "MatMul"], [1, None, 0, 0, 0, 0], ), self.match_parent_path( start_node, ["MatMul", "Mul", "MatMul", "Mul", "Div", "Add"], [None, 0, None, 0, None, 0], ), self.match_parent_path( start_node, ["MatMul", "Mul", "MatMul", "SimplifiedLayerNormalization", "Add"], [None, 0, None, 0, 0], ), ] if qkv_nodes_list is not None and any(qkv_nodes_list): start_nodes.append(start_node) # can't find qkv nodes with above patterns, use Softmax nodes to split model if len(start_nodes) == 0: for node in self.model.graph.node: if node.op_type == "Softmax": start_nodes.append(node) return start_nodes def find_split_nodes(self): """Find split nodes for layer-wise quantization.""" self.remove_unused_nodes() split_nodes = self.find_split_node_for_layer_wise_quantization() return split_nodes def _infer_tensor_dtype(self): """Infer the elem_type of tensors.""" initializers = dict([(i.name, i.data_type) for i in self.model.graph.initializer]) inputs = dict([(i.name, i.type.tensor_type.elem_type) for i in self.model.graph.input]) value_info = dict([(i.name, i.type.tensor_type.elem_type) for i in self.model.graph.value_info]) outputs = dict([(i.name, i.type.tensor_type.elem_type) for i in self.model.graph.output]) for node in self.model.graph.node: if node.output[0] in value_info: continue elem_type = None if node.op_type in ["And", "Equal", "Greater", "GreaterOrEqual", "Less", "LessOrEqual", "Or", "Xor"]: elem_type = onnx.TensorProto.BOOL elif node.op_type in ["ArgMax", "ArgMin", "NonZero", "Shape"]: elem_type = onnx.TensorProto.INT64 elif node.op_type == "Cast" and len(node.attribute) > 0: elem_type = node.attribute[0].i elif node.op_type in ["Constant", "ConstantOfShape"] and len(node.attribute) > 0: elem_type = node.attribute[0].t.data_type elif len(node.input) >= 2: for inp in node.input[:2]: if inp in initializers and initializers[inp] != onnx.TensorProto.INT64: elem_type = initializers[inp] break # output elem_type aligns with input if elem_type is None and len(node.input) > 0: inp = node.input[0] if inp in value_info: elem_type = value_info[inp] elif inp in inputs: elem_type = inputs[inp] elif inp in outputs: elem_type = outputs[inp] if elem_type is not None: if node.op_type in ["Split", "Slice"]: for out in node.output: value_info.update({out: elem_type}) else: value_info.update({node.output[0]: elem_type}) return value_info def _build_input_output_tensor(self, tensor_name, value_info): if tensor_name in self.input(): return self.model.graph.input[self.input().index(tensor_name)] if tensor_name in self.output(): return self.model.graph.output[self.output().index(tensor_name)] tensor_type = value_info.get(tensor_name, onnx.TensorProto.FLOAT) return onnx.helper.make_tensor_value_info(tensor_name, tensor_type, None) def split_model_with_node( self, split_node_name, path_of_model_to_split, save_both_split_models=True, save_path=None ): """Split model into two parts at a given node. Args: split_node_name (str): name of the node where the model is split at> path_of_model_to_split (str): path of model to be split. save_both_split_models (bool): whether to save the two split models. False means only save the first split model. True means save both the two split models. Default id True. save_path (str): path to save split models. None means using self.model_path Returns: tuple: the first split model, the second split model """ # origin model : ... -> node_1 -> split_node -> node_2 -> ... # split model 1: ... -> node_1 -> split_node # split model 2: node_2 -> ... # infer elem_type of tensors to make sure layer-wise quant run successfully value_info = self._infer_tensor_dtype() split_model_part_1 = onnx.ModelProto() split_model_part_1.CopyFrom(self.model) split_model_part_1.graph.ClearField("node") split_model_part_2 = onnx.ModelProto() split_model_part_2.CopyFrom(self.model) split_model_part_2.graph.ClearField("node") split_node = None nodes = [] for node in self.model.graph.node: nodes.append(node) if node.name == split_node_name: split_node = node break assert len(split_node.output) == 1, ( "Only support split at node with 1 output tensor, while " "current split node {} has {} output tensors".format(split_node_name, len(split_node.output)) ) split_tensor_name = split_node.output[0] split_tensor = self._build_input_output_tensor(split_tensor_name, value_info) split_model_part_1.graph.node.extend(nodes) split_model_part_1.graph.output.append(split_tensor) split_model_part_1 = ONNXModel(split_model_part_1, ignore_warning=True) # remove isolated graphs which are not related to the split_node output_name_to_node = split_model_part_1.output_name_to_node() valid_nodes = [split_node] while len(valid_nodes) > 0: node = valid_nodes.pop(0) for inp in node.input: if inp in output_name_to_node: valid_nodes.append(output_name_to_node[inp]) if node in nodes: nodes.remove(node) split_model_part_1.remove_nodes(nodes) for node in self.model.graph.node: if node not in split_model_part_1.nodes(): split_model_part_2.graph.node.append(node) split_model_part_2.graph.input.append(split_tensor) split_model_part_2 = ONNXModel(split_model_part_2, ignore_warning=True) # remove unused input & output split_model_part_1._remove_unused_input_output() split_model_part_2._remove_unused_input_output() insert_output_for_model_1 = [] insert_input_for_model_2 = [] for output in split_model_part_1._output_name_to_node.keys(): if output in split_model_part_2._input_name_to_nodes.keys(): output_tensor = self._build_input_output_tensor(output, value_info) if output_tensor not in split_model_part_1.model.graph.output: insert_output_for_model_1.append(output_tensor) if output_tensor not in split_model_part_2.model.graph.input: insert_input_for_model_2.append(output_tensor) # insert model 1 output for output in insert_output_for_model_1: split_model_part_1.model.graph.output.append(output) # insert model 2 input for input in insert_input_for_model_2: split_model_part_2.model.graph.input.append(input) # remove unused init split_model_part_1.remove_unused_init() split_model_part_2.remove_unused_init() split_model_part_1.update() split_model_part_2.update() dir_of_model_to_split = os.path.dirname(path_of_model_to_split) split_model_part_1.load_model_initializer_by_tensor(dir_of_model_to_split) split_model_part_1_path = ( os.path.join(save_path, "split_model_part_1.onnx") if save_path is not None else os.path.join(dir_of_model_to_split, "split_model_part_1.onnx") ) split_model_part_1.model_path = split_model_part_1_path split_model_part_1._save_split_model(split_model_part_1_path) split_model_part_1.check_is_large_model() logger.debug("save split model part 1 to {} for layer wise quantization".format(split_model_part_1_path)) if save_both_split_models: split_model_part_2.load_model_initializer_by_tensor(dir_of_model_to_split) split_model_part_2_path = ( os.path.join(save_path, "split_model_part_2.onnx") if save_path is not None else os.path.join(dir_of_model_to_split, "split_model_part_2.onnx") ) split_model_part_2.model_path = split_model_part_2_path split_model_part_2._save_split_model(split_model_part_2_path) split_model_part_2.check_is_large_model() logger.debug("save split model part 2 to {} for layer wise quantization".format(split_model_part_2_path)) return split_model_part_1, split_model_part_2 else: return split_model_part_1, split_model_part_2 def _save_split_model(self, save_path): """Save split model as external data for layer wise quantization. Args: save_path (str): the path to save the split model """ if os.path.exists(save_path + "_data"): os.remove(save_path + "_data") self._model_path = save_path onnx.save_model( self.model, save_path, save_as_external_data=True, all_tensors_to_one_file=True, location=os.path.basename(save_path) + "_data", size_threshold=1024, convert_attribute=False, ) def _remove_unused_input_output(self): """Remove unused input & output for split model.""" remove_outputs = [] remove_inputs = [] input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() for output in self.model.graph.output: if output.name not in output_name_to_node.keys(): remove_outputs.append(output) for input in self.model.graph.input: if input.name not in input_name_to_nodes.keys(): remove_inputs.append(input) for output in remove_outputs: self.model.graph.output.remove(output) for input in remove_inputs: self.model.graph.input.remove(input) def remove_unused_init(self): """Remove unused init.""" remov_inits = [] if len(self._input_name_to_nodes) == 0: self._input_name_to_nodes = self.input_name_to_nodes() for init in self.model.graph.initializer: if init.name not in self._input_name_to_nodes.keys(): remov_inits.append(init) self.remove_initializers(remov_inits) def load_model_initializer_by_tensor(self, data_path=None): """Load model initializer by tensor. Args: data_path (str, optional): the directory of saved initializer. Defaults to None. """ if data_path is None: data_path = os.path.dirname(self._model_path) for init in self.model.graph.initializer: if init.HasField("data_location") and init.data_location == onnx.TensorProto.EXTERNAL: onnx.external_data_helper.load_external_data_for_tensor(init, data_path) def write_external_data_to_new_location(self, external_data_location="external.data", overwrite=False): """Write external data of merged quantized model to new location to save memory. Args: external_data_location (str, optional): external data location of merged quantized model. Defaults to "external.data". overwrite (bool, optional): if True, remove existed externa data. Defaults to False. """ if overwrite and os.path.exists(os.path.join(os.path.dirname(self._model_path), external_data_location)): os.remove(os.path.join(os.path.dirname(self._model_path), external_data_location)) self.load_model_initializer_by_tensor() onnx.external_data_helper.convert_model_to_external_data(self.model, location=external_data_location) # TODO : if init is already saved, skip write it onnx.external_data_helper.write_external_data_tensors(self.model, filepath=os.path.dirname(self._model_path)) def merge_split_models(self, to_merge_model): """Merge two split model into final model.""" to_merge_model.write_external_data_to_new_location() self.add_nodes([node for node in to_merge_model.nodes()]) self.add_initializers([init for init in to_merge_model.initializer()]) self.update() # add new output for output in to_merge_model.graph().output: if output.name not in self.output(): self.model.graph.output.append(output) # remove unused output remove_output = [] for output in self.model.graph.output: if output.name in to_merge_model.input(): remove_output.append(output) for output in remove_output: self.model.graph.output.remove(output) # add new input for input in to_merge_model.graph().input: if ( input.name not in self.input() and input.name not in self.output() and input.name not in self._output_name_to_node.keys() ): self.model.graph.input.append(input) def re_org_output(self, origin_output): """Re-org output of merged model for layer-wise quantization.""" outputs = {} tmp_remove = [] for output in self.model.graph.output: outputs[output.name] = output tmp_remove.append(output) for output in tmp_remove: self.model.graph.output.remove(output) for out_name in origin_output: self.model.graph.output.append(outputs[out_name]) ================================================ FILE: onnx_neural_compressor/utility.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import importlib import logging import os import pathlib import subprocess import time import cpuinfo import numpy as np import onnx import onnxruntime as ort import prettytable as pt import psutil from onnx_neural_compressor import constants, logger from typing import Callable, Dict, List, Tuple, Union # isort: skip # Dictionary to store a mapping between algorithm names and corresponding algo implementation(function) algos_mapping: Dict[str, Callable] = {} ####################################################### #### Options ####################################################### def check_value(name, src, supported_type, supported_value=[]): """Check if the given object is the given supported type and in the given supported value. Example:: from onnx_neural_compressor import utility def datatype(self, datatype): if utility.check_value("datatype", datatype, list, ["fp32", "bf16", "uint8", "int8"]): self._datatype = datatype """ if isinstance(src, list) and any([not isinstance(i, supported_type) for i in src]): assert False, "Type of '{}' items should be {} but not {}".format( name, str(supported_type), [type(i) for i in src] ) elif not isinstance(src, list) and not isinstance(src, supported_type): assert False, "Type of '{}' should be {} but not {}".format(name, str(supported_type), type(src)) if len(supported_value) > 0: if isinstance(src, str) and src not in supported_value: assert False, "'{}' is not in supported '{}': {}. Skip setting it.".format(src, name, str(supported_value)) elif ( isinstance(src, list) and all([isinstance(i, str) for i in src]) and any([i not in supported_value for i in src]) ): assert False, "{} is not in supported '{}': {}. Skip setting it.".format(src, name, str(supported_value)) return True class Options: """Option Class for configs. This class is used for configuring global variables. The global variable options is created with this class. If you want to change global variables, you should use functions from onnx_neural_compressor.utility.py: set_random_seed(seed: int) Args: random_seed(int): Random seed used in neural compressor. Default value is 1978. Example:: from onnx_neural_compressor import set_random_seed set_random_seed(2022) """ def __init__(self, random_seed=1978): """Init an Option object.""" self.random_seed = random_seed @property def random_seed(self): """Get random seed.""" return self._random_seed @random_seed.setter def random_seed(self, random_seed): """Set random seed.""" if check_value("random_seed", random_seed, int): self._random_seed = random_seed options = Options() def singleton(cls): """Singleton decorator.""" instances = {} def _singleton(*args, **kw): """Create a singleton object.""" if cls not in instances: instances[cls] = cls(*args, **kw) return instances[cls] return _singleton class Statistics: """The statistics printer.""" def __init__(self, data, header, field_names, output_handle=logger.info): """Init a Statistics object. Args: data: The statistics data header: The table header field_names: The field names output_handle: The output logging method """ self.field_names = field_names self.header = header self.data = data self.output_handle = output_handle self.tb = pt.PrettyTable(min_table_width=40) def print_stat(self): """Print the statistics.""" valid_field_names = [] for index, value in enumerate(self.field_names): if index < 2: valid_field_names.append(value) continue if any(i[index] for i in self.data): valid_field_names.append(value) self.tb.field_names = valid_field_names for i in self.data: tmp_data = [] for index, value in enumerate(i): if self.field_names[index] in valid_field_names: tmp_data.append(value) if any(tmp_data[1:]): self.tb.add_row(tmp_data) lines = self.tb.get_string().split("\n") self.output_handle("|" + self.header.center(len(lines[0]) - 2, "*") + "|") for i in lines: self.output_handle(i) class LazyImport(object): """Lazy import python module till use.""" def __init__(self, module_name): """Init LazyImport object. Args: module_name (string): The name of module imported later """ self.module_name = module_name self.module = None def __getattr__(self, name): """Get the attributes of the module by name.""" try: self.module = importlib.import_module(self.module_name) mod = getattr(self.module, name) except: spec = importlib.util.find_spec(str(self.module_name + "." + name)) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod def __call__(self, *args, **kwargs): """Call the function in that module.""" function_name = self.module_name.split(".")[-1] module_name = self.module_name.split(f".{function_name}")[0] self.module = importlib.import_module(module_name) function = getattr(self.module, function_name) return function(*args, **kwargs) @singleton class CpuInfo(object): """CPU info collection.""" def __init__(self): """Get whether the cpu numerical format is bf16, the number of sockets, cores and cores per socket.""" self._bf16 = False self._vnni = False info = cpuinfo.get_cpu_info() if "arch" in info and "X86" in info["arch"]: cpuid = cpuinfo.CPUID() max_extension_support = cpuid.get_max_extension_support() if max_extension_support >= 7: ecx = cpuid._run_asm( b"\x31\xC9", # xor ecx, ecx b"\xB8\x07\x00\x00\x00" b"\x0f\xa2" b"\x89\xC8" b"\xC3", # mov eax, 7 # cpuid # mov ax, cx # ret ) self._vnni = bool(ecx & (1 << 11)) eax = cpuid._run_asm( b"\xB9\x01\x00\x00\x00", # mov ecx, 1 b"\xB8\x07\x00\x00\x00" b"\x0f\xa2" b"\xC3", # mov eax, 7 # cpuid # ret ) self._bf16 = bool(eax & (1 << 5)) # TODO: The implementation will be refined in the future. # https://github.com/intel/neural-compressor/tree/detect_sockets if "arch" in info and "ARM" in info["arch"]: # pragma: no cover self._sockets = 1 else: self._sockets = self.get_number_of_sockets() self._cores = psutil.cpu_count(logical=False) self._cores_per_socket = int(self._cores / self._sockets) @property def bf16(self): """Get whether it is bf16.""" return self._bf16 @property def vnni(self): """Get whether it is vnni.""" return self._vnni @property def cores_per_socket(self): """Get the cores per socket.""" return self._cores_per_socket def get_number_of_sockets(self) -> int: """Get number of sockets in platform.""" cmd = "cat /proc/cpuinfo | grep 'physical id' | sort -u | wc -l" if psutil.WINDOWS: cmd = r'wmic cpu get DeviceID | C:\Windows\System32\find.exe /C "CPU"' elif psutil.MACOS: # pragma: no cover cmd = "sysctl -n machdep.cpu.core_count" with subprocess.Popen( args=cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=False, ) as proc: proc.wait() if proc.stdout: for line in proc.stdout: return int(line.decode("utf-8", errors="ignore").strip()) return 0 def set_random_seed(seed: int): """Set the random seed in config.""" options.random_seed = seed def simple_progress_bar(total, i): """Progress bar for cases where tqdm can't be used.""" progress = i / total bar_length = 20 bar = "#" * int(bar_length * progress) spaces = " " * (bar_length - len(bar)) percentage = progress * 100 print(f"\rProgress: [{bar}{spaces}] {percentage:.2f}%", end="") def register_algo(name): """Decorator function to register algorithms in the algos_mapping dictionary. Usage example: @register_algo(name=example_algo) def example_algo(model: Union[onnx.ModelProto, pathlib.Path, str], quant_config: RTNConfig) -> onnx.ModelProto: ... Args: name (str): The name under which the algorithm function will be registered. Returns: decorator: The decorator function to be used with algorithm functions. """ def decorator(algo_func): algos_mapping[name] = algo_func return algo_func return decorator def auto_detect_ep(): eps = ort.get_available_providers() if "DnnlExecutionProvider" in eps: return "DnnlExecutionProvider" elif "DmlExecutionProvider" in eps: return "DmlExecutionProvider" elif "CUDAExecutionProvider" in eps: return "CUDAExecutionProvider" else: return "CPUExecutionProvider" def trt_env_setup(model): """Set environment variable for Tensorrt Execution Provider.""" is_int8 = False for node in model.graph.node: if node.op_type in ["QuantizeLinear", "DequantizeLinear"]: is_int8 = True break if is_int8: os.environ["ORT_TENSORRT_INT8_ENABLE"] = "1" else: os.environ["ORT_TENSORRT_INT8_ENABLE"] = "0" ================================================ FILE: onnx_neural_compressor/version.py ================================================ # Copyright (c) 2021 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Neural Compressor: An open-source Python library supporting popular model compression techniques for ONNX models.""" __version__ = "1.0" ================================================ FILE: onnx_neural_compressor/algorithms/__init__.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ================================================ FILE: onnx_neural_compressor/algorithms/utility.py ================================================ # Copyright (c) 2023 MIT HAN Lab # This source code is licensed under the MIT license # # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import re import struct import sys from importlib import util import numpy as np from packaging import version from onnx_neural_compressor import constants, utility if sys.version_info < (3, 11) and util.find_spec("onnxruntime_extensions"): # pragma: no cover import onnxruntime_extensions onnx = utility.LazyImport("onnx") ort = utility.LazyImport("onnxruntime") __producer__ = "onnx.quantize" __version__ = "0.1.0" onnx_domain = "ai.onnx" ms_domain = "com.microsoft" QUANT_OP_NAME_SUFFIX = "_quant" def attribute_to_kwarg(attribute): """Convert attribute to kwarg format for use with onnx.helper.make_node.""" attribute_mapping = { 1: attribute.f, 2: attribute.i, 3: attribute.s, 4: attribute.t, 5: attribute.g, 6: attribute.floats, 7: attribute.ints, 8: attribute.strings, 9: attribute.tensors, 10: attribute.graphs, } if attribute.type in attribute_mapping: value = attribute_mapping[attribute.type] else: # pragma: no cover raise ValueError( "attribute {} has no type specified " "or unsupported type {}.".format(attribute.name, attribute.type) ) return {attribute.name: value} ONNX_INT_TYPE_RANGE = { onnx.TensorProto.UINT8: (0, 255), onnx.TensorProto.INT8: (-128, 127), } ONNX_INT_TYPE_SYMMETRIC_RANGE = { onnx.TensorProto.INT8: (-127, 127), } ONNX_INT_TYPE_REDUCED_RANGE = { onnx.TensorProto.UINT8: (0, 127), onnx.TensorProto.INT8: (-64, 64), } ONNX_STR_TYPE_RANGE = { "int1": (-1, 0), "int2": (-2, 1), "int3": (-4, 3), "int4": (-8, 7), # onnx >= 1.16.0 defines TensorProto.INT4 "int5": (-16, 15), "int6": (-32, 31), "int7": (-64, 63), "int8": (-128, 127), "uint1": (0, 1), "uint2": (0, 3), "uint3": (0, 7), "uint4": (0, 15), # onnx >= 1.16.0 defines TensorProto.UINT4 "uint5": (0, 31), "uint6": (0, 63), "uint7": (0, 127), "uint8": (0, 255), } ONNX_TENSOR_TYPE = { "bfloat16": getattr(onnx.TensorProto, "BFLOAT16", 16), "float32": getattr(onnx.TensorProto, "FLOAT", 1), "float16": getattr(onnx.TensorProto, "FLOAT16", 10), "int4": getattr(onnx.TensorProto, "INT4", 22), "uint4": getattr(onnx.TensorProto, "UNT4", 21), "int8": getattr(onnx.TensorProto, "INT8", 3), "uint8": getattr(onnx.TensorProto, "UINT8", 2), } def _qType_to_np_type(qType): if isinstance(qType, int): return onnx.helper.tensor_dtype_to_np_dtype(qType) elif isinstance(qType, str) and "uint" in qType: return np.dtype("uint8") else: return np.dtype("int8") def find_by_name(name, item_list): """Helper function to find item by name in a list.""" items = [] for item in item_list: assert hasattr(item, "name"), "{} should have a 'name' attribute defined".format(item) if item.name == name: items.append(item) if len(items) > 0: return items[0] else: return None def get_qmin_qmax_for_qType(qType, reduce_range=False, sym=False): # noqa: N802 """Get qmin, qmax for qType. Args: qType (int or str): int for onnx defined type, str for onnx not defined type reduce_range (bool, optional): whether use 7 bit for 8bit quantization sym (bool, optional): quantization scheme. Defaults to False. """ if qType == onnx.TensorProto.FLOAT8E4M3FN: raise NotImplementedError("This function is not implemented for float 8 as not needed.") qrange = None if isinstance(qType, str): qrange = ONNX_STR_TYPE_RANGE.get(qType) elif reduce_range: qrange = ONNX_INT_TYPE_REDUCED_RANGE.get(qType) elif sym and qType in ONNX_INT_TYPE_SYMMETRIC_RANGE: qrange = ONNX_INT_TYPE_SYMMETRIC_RANGE[qType] else: qrange = ONNX_INT_TYPE_RANGE.get(qType) if not qrange: raise ValueError(f"Unexpected data type {qType} requested.") return qrange def quantize_nparray(dtype, arr, scale, zero_point, low=None, high=None): """Quantize numpy array.""" q_weight = np.empty_like(np.asarray(arr), dtype=np.asarray(scale).dtype) np.divide(arr, scale, out=q_weight) np.add(q_weight, zero_point, out=q_weight) np.round(q_weight, out=q_weight) if low is not None and high is not None: np.clip(q_weight, low, high, out=q_weight) return q_weight.astype(dtype) def quantize_data_per_channel(data, axis, qType, sym, reduce_range=False): """Quantize tensor per-channel.""" quantize_range = get_qmin_qmax_for_qType(qType, reduce_range, sym) rmin = None rmax = None for i in range(len(data.shape)): if i != axis: rmin = np.min(data, axis=i, keepdims=True) if rmin is None else np.min(rmin, axis=i, keepdims=True) rmax = np.max(data, axis=i, keepdims=True) if rmax is None else np.max(rmax, axis=i, keepdims=True) rmin = np.minimum(rmin, 0) rmax = np.maximum(rmax, 0) scale, zero_point = calculate_scale_zp(rmin, rmax, qType, sym, reduce_range) dtype = _qType_to_np_type(qType) quantized_data = quantize_nparray(dtype, data, scale, zero_point, low=quantize_range[0], high=quantize_range[1]) return rmin.reshape(-1, 1), rmax.reshape(-1, 1), zero_point.reshape(-1, 1), scale.reshape(-1, 1), quantized_data def dequantize_data_with_scale_zero(tensor_value, scale_value, zo_value): """Dequantize tensor with scale and zero point.""" return (tensor_value.astype(scale_value.dtype) - zo_value.astype(scale_value.dtype)) * scale_value def dequantize_data(tensor_value, scale_value, zo_value, axis=0): """Dequantize tensor.""" if not isinstance(scale_value, np.ndarray): return dequantize_data_with_scale_zero(tensor_value, scale_value, zo_value) else: channel_count = tensor_value.shape[axis] # TBD, default from axis 0 new_per_channel_tensor_values = [] for i in range(channel_count): per_channel_tensor_value = tensor_value.take(i, axis) per_channel_scale_value = scale_value.take(i) per_channel_zero_value = zo_value.take(i) new_per_channel_tensor_values.append( dequantize_data_with_scale_zero( per_channel_tensor_value, per_channel_scale_value, per_channel_zero_value ) ) # combine per_channel_data into one reshape_dims = list(tensor_value.shape) # deep copy reshape_dims[axis] = 1 # only one per channel for reshape new_tensor_value = new_per_channel_tensor_values[0].reshape(reshape_dims) for i in range(1, channel_count): new_per_channel_tensor_value = new_per_channel_tensor_values[i].reshape(reshape_dims) new_tensor_value = np.concatenate((new_tensor_value, new_per_channel_tensor_value), axis) return new_tensor_value def calculate_scale_zp(rmin, rmax, qType, sym, reduce_range=False): """Calculate scale and zero point.""" qmin, qmax = get_qmin_qmax_for_qType(qType, reduce_range, sym) dtype = _qType_to_np_type(qType) if isinstance(rmax, np.ndarray): if sym: max_range = np.maximum(abs(rmin), abs(rmax)) rmin = -max_range rmax = max_range scale = (rmax - rmin) / (qmax - qmin) scale[abs(scale) < np.finfo(rmax.dtype).tiny] = 1 zero_point = ( np.multiply(np.ones(rmax.shape), np.round((qmax + qmin) / 2.0)).astype(dtype) if sym else np.round(qmin - rmin / scale).astype(dtype) ) else: if sym: max_range = max(abs(rmin), abs(rmax)) scale = (float(max_range) * 2) / (qmax - qmin) if max_range > 0 else 1 else: scale = (float(rmax) - float(rmin)) / (qmax - qmin) if rmin != rmax else 1 zero_point = np.round((qmax + qmin) / 2.0).astype(dtype) if sym else np.round(qmin - rmin / scale).astype(dtype) return np.float32(scale), zero_point def quantize_data(data, qType, sym, reduce_range=False, ratio=1.0, axis=None): """Quantize data. To pack weights, we compute a linear transformation - when data type == uint8 mode, from [rmin, rmax] -> [0, 2^{b-1}] and - when data type == int8, from [-m , m] -> [-(2^{b-1}-1), 2^{b-1}-1] where m = max(abs(rmin), abs(rmax)) and add necessary intermediate nodes to transform quantized weight to full weight using the equation r = S(q-z), where r: real original value q: quantized value S: scale z: zero point Args: data (array): data to quantize qType (int): data type to quantize to. Supported types UINT8 and INT8 sym (bool): whether use sym quantization. reduce_range (bool): whether use 7 bit or not. Defaults to False ratio (float, optional): percentile of clip. Defaults to 1.0 axis (int, optional): process data along a specific axis. Default is None (process the whole data) """ quantize_range = get_qmin_qmax_for_qType(qType, reduce_range, sym) rmin = np.min(np.min(data), 0) if axis is None else np.min(data, axis=axis, keepdims=True) rmax = np.max(np.max(data), 0) if axis is None else np.max(data, axis=axis, keepdims=True) rmin *= ratio rmax *= ratio scale, zero_point = calculate_scale_zp(rmin, rmax, qType, sym, reduce_range) dtype = _qType_to_np_type(qType) quantized_data = quantize_nparray(dtype, data, scale, zero_point, low=quantize_range[0], high=quantize_range[1]) return rmin, rmax, zero_point, scale, quantized_data def qdq_data(data, qType, sym, reduce_range=False, ratio=1.0, axis=None): _, _, zero_point, scale, quantized_data = quantize_data(data, qType, sym, reduce_range, ratio, axis) return scale * (quantized_data - zero_point) def is_B_transposed(node): """Whether inuput B is transposed.""" transB = [attr for attr in node.attribute if attr.name == "transB"] if len(transB): return 0 < onnx.helper.get_attribute_value(transB[0]) return False def is_quantizable_type(data_type): return data_type in [onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16, onnx.TensorProto.BFLOAT16] def _get_blob_size(group_size, has_zp): # pragma: no cover """Get blob_size. Args: group_size (int): how many elements share one scale/zp has_zp (bool): whether zero_point is None """ if version.Version(ort.__version__) > constants.ONNXRT1161_VERSION: blob_size = group_size // 2 elif has_zp: blob_size = group_size // 2 + 4 + 1 else: blob_size = group_size // 2 + 4 return blob_size def make_weight_only_dequant_node( node: onnx.NodeProto, weight_shape: tuple, block_size: int, num_bits: int, dtype: str, q_weight: np.array, scale: np.array, zero_point: np.array, axis: int = 1, ): """Build DequantizeLinear node. Args: node: original matmul node weight_shape (tuple): original weight shape block_size (int): how many elements share one scale/zp num_bits (int): num_bits dtype (str): use uint or int q_weight (array): quantized weight scale (array): scale zero_point (array): zero point axis (int): the axis of the dequantizing dimension of the input tensor Returns: weight_only_dequant_node: DequantizeLinear node for weight dequantization new_inits: initializers of the new node """ new_inits = [] input_names = [] kwargs = {"block_size": block_size, "axis": axis} q_weight = q_weight.reshape((weight_shape[-1], -1)).T if num_bits == 4: q_weight = ((q_weight[:, ::2] & 0xF | q_weight[:, 1::2] << 4) & 0xFF).astype("uint8") qtype = ONNX_TENSOR_TYPE.get(dtype + str(num_bits), None) if qtype is None: raise ValueError( "Unsupported qtype {}, only support {}".format(dtype + str(num_bits), list(ONNX_TENSOR_TYPE.keys())) ) q_weight_tensor = onnx.helper.make_tensor( name=node.input[1] + "_Q{}G{}".format(str(num_bits), str(block_size)), data_type=qtype, dims=weight_shape, vals=q_weight.flatten().tobytes(), raw=True, ) new_inits.append(q_weight_tensor) input_names.append(q_weight_tensor.name) scale = scale.reshape((weight_shape[-1], -1)).T scale_tensor = onnx.helper.make_tensor( name=node.input[1] + "_scale", data_type=onnx.helper.np_dtype_to_tensor_dtype(scale.dtype), dims=scale.shape, vals=scale.tobytes(), raw=True, ) input_names.append(scale_tensor.name) new_inits.append(scale_tensor) # build zero_point tensor zero_point = zero_point.reshape((weight_shape[-1], -1)).T if num_bits == 4: zero_point = ((zero_point[:, ::2] & 0xF | zero_point[:, 1::2] << 4) & 0xFF).astype("uint8") zp_tensor = onnx.helper.make_tensor( name=node.input[1] + "_zp", data_type=qtype, dims=scale.shape, vals=zero_point.flatten().tobytes(), raw=True, ) input_names.append(zp_tensor.name) new_inits.append(zp_tensor) dequant_node = onnx.helper.make_node( "DequantizeLinear", inputs=input_names, outputs=[q_weight_tensor.name + "_dequant"], name=node.name + "_woq_dequant", **kwargs, ) node.input[1] = dequant_node.output[0] return dequant_node, new_inits def make_matmul_weight_only_node( node: onnx.NodeProto, weight_shape: tuple, num_bits: int, group_size: int, k_blocks: int, q_weight: np.array, scale: np.array, zero_point: np.array, accuracy_level: int = 0, ): """Build MatMulFpQ4/MatMulNBits node. Args: node (onnx.NodeProto): original matmul node weight_shape (tuple): original weight shape num_bits (int): number of bits used to represent weights. group_size (int): how many elements share one scale/zp k_blocks (int): block number q_weight (np.array): quantized weight scale (np.array): scale zero_point (np.array): zero point accuracy_level (int, optional): accuracy level. Support 0 (unset), 1(fp32 compute type of jblas kernel), 2 (fp16 compute type of jblas kernel), 3 (bf16 compute type of jblas kernel), 4 (int8 compute type of jblas kernel) Defaults to 0. Returns: matmul_weight_only_node: MatMulFpQ4 or MatMulNBits node new_inits: initializers of the new node """ blob_size = _get_blob_size(group_size, zero_point is not None) packed = np.zeros((q_weight.shape[0], blob_size), dtype="uint8") q_weight_name = node.input[1] + "_Q{}G{}".format(str(num_bits), str(group_size)) input_names = [node.input[0], q_weight_name] new_inits = [] kwargs = {} if version.Version(ort.__version__) > constants.ONNXRT1161_VERSION: op_type = "MatMulNBits" # pack quantized weight q_weight_pairs = q_weight[:, ::2] | q_weight[:, 1::2] << 4 packed[:, :] = q_weight_pairs[:, :blob_size] packed = np.reshape(packed, (-1, k_blocks, blob_size)) # build scale tensor scale = np.reshape(scale, (-1, k_blocks)) scale_tensor = onnx.helper.make_tensor( name=node.input[1] + "_scale", data_type=onnx.helper.np_dtype_to_tensor_dtype(scale.dtype), dims=scale.shape, vals=scale.tobytes(), raw=True, ) input_names.append(scale_tensor.name) new_inits.append(scale_tensor) # build zero_point tensor if zero_point is not None: if num_bits > 4: packed_zp = np.reshape(zero_point, (1, -1)).astype("uint8") else: packed_zp = np.full((zero_point.shape[0] + 1) // 2, 136, dtype="uint8") # create an index array idx = np.arange(zero_point.shape[0] // k_blocks * k_blocks).reshape(-1) # separate odd and even indices even_idx = idx[::2] odd_idx = idx[1::2] # vectorized operation for even and odd indices packed_zp[even_idx // 2] = (packed_zp[even_idx // 2] & 0xF0) | zero_point[even_idx].ravel() packed_zp[odd_idx // 2] = (packed_zp[odd_idx // 2] & 0x0F) | (zero_point[odd_idx].ravel() << 4) zp_tensor = onnx.helper.make_tensor( name=node.input[1] + "_zp", data_type=2, dims=packed_zp.shape, vals=packed_zp.tobytes(), raw=True ) input_names.append(zp_tensor.name) new_inits.append(zp_tensor) # set kwargs kwargs["K"] = weight_shape[0] kwargs["N"] = weight_shape[1] kwargs["bits"] = num_bits kwargs["block_size"] = group_size if accuracy_level > 0: # require onnxruntime > 1.16.3 kwargs["accuracy_level"] = accuracy_level else: # pragma: no cover offset = 5 if zero_point is not None else 4 op_type = "MatMulFpQ4" # pack quantized weight for i in range(q_weight.shape[0]): bf = struct.pack("f", scale[i]) packed[i][0] = bf[0] packed[i][1] = bf[1] packed[i][2] = bf[2] packed[i][3] = bf[3] if zero_point is not None: packed[i][4] = zero_point[i] packed[i][offset:] = np.bitwise_or( q_weight[i][: group_size // 2], np.left_shift(q_weight[i][group_size // 2 :], num_bits) ) packed = packed.reshape(-1) # build shape tensor shape_tensor = onnx.helper.make_tensor( name=node.input[1] + "_shape", data_type=7, dims=(2,), vals=np.array(weight_shape, dtype="int64") ) new_inits.append(shape_tensor) input_names.append(shape_tensor.name) # set kwargs kwargs["blk_quant_type"] = 1 if zero_point is not None else 0 q_weight_tensor = onnx.helper.make_tensor( name=q_weight_name, data_type=2, dims=packed.shape, vals=packed.tobytes(), raw=True, ) new_inits.append(q_weight_tensor) matmul_weight_only_node = onnx.helper.make_node( op_type, inputs=input_names, outputs=node.output, name=node.name + "_Q" + str(num_bits) if node.name else "_Q" + str(num_bits), domain="com.microsoft", **kwargs, ) return matmul_weight_only_node, new_inits def quant_matmul_weight_only( node, weight, dtype, num_bits, sym, group_size, ratio=1, quant_format=None, accuracy_level=0, ): new_nodes = [] new_inits = [] remove_nodes = [] org_w_shape = weight.shape # ic, oc group_size = group_size if group_size != -1 else org_w_shape[0] k_blocks = (org_w_shape[0] - 1) // group_size + 1 weight = pad_tensor(weight, group_size, k_blocks) if quant_format == 1: _, _, zp, scale, q_weight = quantize_data( weight.T.reshape((-1, group_size)), dtype + str(num_bits), sym, ratio=ratio, axis=1, ) dequant_node, inits = make_weight_only_dequant_node( node=node, weight_shape=org_w_shape, num_bits=num_bits, dtype=dtype, q_weight=q_weight, scale=scale.astype(weight.dtype), axis=0, block_size=group_size, zero_point=zp, ) new_nodes.append(dequant_node) new_inits.extend(inits) elif quant_format == 0: _, _, zp, scale, q_weight = quantize_data( weight.T.reshape((-1, group_size)), dtype + str(num_bits), sym, ratio=ratio, axis=1, ) q_matmul_node, inits = make_matmul_weight_only_node( node=node, weight_shape=org_w_shape, num_bits=num_bits, group_size=group_size, k_blocks=k_blocks, q_weight=q_weight, scale=scale.astype(weight.dtype), zero_point=zp if not sym else None, accuracy_level=accuracy_level, ) new_nodes.append(q_matmul_node) new_inits.extend(inits) remove_nodes.append(node) else: q_weight = qdq_data( weight.T.reshape((-1, group_size)), dtype + str(num_bits), sym, ratio=ratio, axis=1, ) q_weight = np.reshape(q_weight, (org_w_shape[1], -1)) q_weight = np.transpose(q_weight) q_weight = q_weight[: org_w_shape[0], :].astype(weight.dtype) q_weight_tensor = onnx.helper.make_tensor( name=node.input[1] + "_Q{}G{}".format(str(num_bits), str(group_size)), data_type=onnx.helper.np_dtype_to_tensor_dtype(q_weight.dtype), dims=weight.shape, vals=q_weight.tobytes(), raw=True, ) node.input[1] = q_weight_tensor.name new_inits.append(q_weight_tensor) return new_nodes, new_inits, remove_nodes def prepare_inputs(model, data_reader, providers): """Prepare inputs for weight only quantization. Args: model (ModelProto or onnx_model.ONNXModel): onnx model. data_reader (CalibrationDataReader): a calibration data reader. providers (list): providers to use. Returns: inputs: prepared inputs. so: session options """ so = ort.SessionOptions() if sys.version_info < (3, 11) and util.find_spec("onnxruntime_extensions"): # pragma: no cover so.register_custom_ops_library(onnxruntime_extensions.get_library_path()) if model.is_large_model: onnx.save_model( model.model, model.model_path + "_augment.onnx", save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) inputs_list = [] while True: inputs = data_reader.get_next() if not inputs: break inputs_list.append(inputs) return inputs_list, so def pad_tensor(weight, group_size, k_blocks): """Pad tensor rowi so that it can be is divisible by group_size. Args: weight (array): weight group_size (int): how many elements share one scale/zp k_blocks (int): the number of block Returns: weight: paded weight """ if group_size == -1: return weight org_w_shape = weight.shape padded_rows = k_blocks * group_size pad_len = padded_rows - org_w_shape[0] if pad_len > 0: weight = np.pad(weight, ((0, pad_len), (0, 0)), "constant") return weight def dump_woq_stats(model, quantize_config, white_list=["MatMul"]): res = {} dtype_set = set() for node in model.graph.node: if node.op_type in ["MatMulFpQ4", "MatMulNBits"]: optype = "MatMul" else: optype = node.op_type if optype not in white_list and optype != "DequantizeLinear": continue if optype not in res: res[optype] = {} dtype = "FP32" for inp in node.input: if re.match("^.*_Q\d*G\d*", inp): Q_position = re.search("_Q\d*", inp) full_position = re.search("_Q\d*G\d*", inp) dtype = "A32W{}G{}".format( inp[Q_position.start() + 2 : Q_position.end()], inp[Q_position.end() + 1 : full_position.end()], ) dtype_set.add(dtype) break res[optype][dtype] = res[optype].get(dtype, 0) + 1 dtype_list = list(dtype_set) for dtype in dtype_list: for optype in res.keys(): if dtype not in res[optype]: res[optype][dtype] = 0 # update stats format for dump. field_names = ["Op Type", "Total"] field_names.extend(dtype_list) output_data = [] for op_type in res.keys(): field_results = [op_type, sum(res[op_type].values())] field_results.extend([res[op_type][dtype] for dtype in dtype_list]) output_data.append(field_results) utility.Statistics(output_data, header="Mixed Precision Statistics", field_names=field_names).print_stat() def get_node_original_name(node) -> str: """Get the original name of the given node.""" node_name: str = node.name # TODO how to handle the unquantized node that has the `_quant` suffix, such as `conv_quant`? if node_name.endswith(QUANT_OP_NAME_SUFFIX): return node_name[: -len(QUANT_OP_NAME_SUFFIX)] else: # For unquantized nodes return node_name def split_shared_bias(model): """Split shared tensor.""" input_name_to_nodes = model.input_name_to_nodes() for input_name, node_list in input_name_to_nodes.items(): if len(node_list) > 1 and input_name in [i.name for i in model.model.graph.initializer]: for node in node_list[1:]: if node.op_type not in ["Conv", "FusedConv"]: continue if len(node.input) > 2 and node.input[2] == input_name: new_input_name = node.input[2] + "_nc_split_" + node.name new_input = onnx.helper.make_tensor( new_input_name, model.get_initializer(input_name).data_type, model.get_initializer(input_name).dims, model.get_initializer(input_name).raw_data, True, ) model.add_initializer(new_input) node.input[2] = new_input_name return model def remove_init_from_model_input(model): """Remove initializer from model input.""" inputs = model.model.graph.input name_to_input = {} for inp in inputs: name_to_input[inp.name] = inp for initializer in model.model.graph.initializer: if initializer.name in name_to_input: inputs.remove(name_to_input[initializer.name]) class QuantizedValue: """Represents a linearly quantized value (input/output/initializer).""" def __init__( self, name, new_quantized_name, scale_name, zero_point_name, axis=None, qType=1, ): """Initialization. Args: name (string): tensor name new_quantized_name (string): quantized tensor name scale_name (string): scale name zero_point_name (string): zero point name axis (int, optional): quantized axis. Defaults to None. qType (int, optional): quantized data type. Defaults to 1 (uint8). """ self.name = name self.q_name = new_quantized_name self.scale_name = scale_name self.zp_name = zero_point_name self.axis = axis self.qType = qType class QuantizedInitializer: """Represents a linearly quantized weight input from ONNX operators.""" def __init__( self, name, initializer, rmins, rmaxs, zero_points, scales, data=[], quantized_data=[], axis=None, qType=1, ): """Initialization. Args: name (string): initializer name initializer (onnx.onnx_ml_pb2.TensorProto): initializer rmins (list): list of min value rmaxs (list): list of max value zero_points (list): list of zero point scales (list): list of scale data (list, optional): array version of the initializer. Defaults to []. quantized_data (list, optional): quantized data. Defaults to []. axis (int, optional): quantized axis. Defaults to None. qType (int, optional): quantized data type. Defaults to 1 (uint8). """ self.name = name self.initializer = initializer # TensorProto initializer in ONNX graph self.rmins = rmins # List of minimum range for each axis self.rmaxs = rmaxs # List of maximum range for each axis # 1D tensor of zero points computed for each axis. scalar if axis is empty self.zero_points = zero_points self.scales = scales # 1D tensor of scales computed for each axis. scalar if axis is empty self.data = data # original data from initializer TensorProto self.quantized_data = quantized_data # weight-packed data from data # Scalar to specify which dimension in the initializer to weight pack. self.axis = axis # If empty, single zero point and scales computed from a single rmin and rmax self.qType = qType def dump_model_op_stats(model, quantize_config, fp32_op_list): qdq_ops = ["QuantizeLinear", "DequantizeLinear", "DynamicQuantizeLinear"] res = {} for op_type in fp32_op_list: res[op_type] = {"INT8": 0, "FP32": 0} for op_type in qdq_ops: res[op_type] = {"INT8": 0, "FP32": 0} for node in model.graph.node: if node.name.endswith("_quant"): if node.op_type.startswith("QLinear"): origin_op_type = node.op_type.split("QLinear")[-1] else: origin_op_type = node.op_type.split("Integer")[0] if origin_op_type in ["QAttention", "QGemm"]: origin_op_type = origin_op_type[1:] elif origin_op_type == "DynamicQuantizeLSTM": origin_op_type = "LSTM" elif origin_op_type == "QEmbedLayerNormalization": origin_op_type = "EmbedLayerNormalization" res[origin_op_type]["INT8"] += 1 elif node.op_type in qdq_ops: res[node.op_type]["INT8"] += 1 elif node.op_type in res: res[node.op_type]["FP32"] += 1 field_names = ["Op Type", "Total", "INT8", "FP32"] output_data = [ [ op_type, sum(res[op_type].values()), res[op_type]["INT8"], res[op_type]["FP32"], ] for op_type in res.keys() ] utility.Statistics(output_data, header="Quantization Statistics", field_names=field_names).print_stat() ================================================ FILE: onnx_neural_compressor/algorithms/layer_wise/__init__.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ================================================ FILE: onnx_neural_compressor/algorithms/layer_wise/core.py ================================================ # Copyright (c) 2023 MIT HAN Lab # This source code is licensed under the MIT license # # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import os import pathlib import tempfile import onnx import onnxruntime as ort from onnx_neural_compressor import data_reader, logger, onnx_model from typing import Callable, List, Union # isort: skip def layer_wise_quant( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], quant_func: Callable, weight_config: dict, data_reader: data_reader.CalibrationDataReader = None, *args, **kwargs ) -> onnx_model.ONNXModel: """Quantize model layer by layer to save memory. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): onnx model. quant_func (Callable): quantization algo function. weight_config (dict): quantization config. data_reader (data_reader.CalibrationDataReader, optional): data_reader for calibration. Defaults to None. Returns: _type_: _description_ """ logger.warning( "Layer-wise quantization requires data_type info for some tensors. " "We will try to infer the data_type automatically if it doesn't exist." "You can use model with symbolic shape inference before layer-wise quantization as well like follows:\n" "import onnxruntime.tools.symbolic_shape_infer as symbolic_shape_infer\n" "model = onnx.load(your_model_path)\n" "out = symbolic_shape_infer.SymbolicShapeInference.infer_shapes(model, auto_merge=True)\n" "onnx.save_model(out, infer_shape_model_path, save_as_external_data=True)\n" ) if not isinstance(model, onnx_model.ONNXModel): model = onnx_model.ONNXModel(model, ignore_warning=True, load_external_data=False) origin_model = copy.deepcopy(model) tmp_file = tempfile.TemporaryDirectory() providers = kwargs.get("providers", ["CPUExecutionProvider"]) # get and check split nodes split_nodes = origin_model.find_split_nodes() if len(split_nodes) == 0: logger.error("Can't find split nodes for layer-wise quantization.") raise ValueError("Fail to run layer-wise quantization.") logger.info( "Will split model into {} parts to do layer-wise quantization".format( len([node.name for node in split_nodes]) + 1 ) ) logger.debug( "Will split model with these nodes for layer-wise quantization: {}".format([node.name for node in split_nodes]) ) split_idx = 1 model_to_split = [origin_model] quantized_model_merged = None require_data_reader = data_reader is not None if require_data_reader: lwq_data_reader = [data_reader] while len(model_to_split) != 0: # prepare model, node and data_reader for current split split_model = model_to_split.pop(0) split_node = split_nodes.pop(0) if require_data_reader: complete_data_reader = lwq_data_reader.pop(0) # if no remaining split nodes, it means this is the last split, and the two split models will be saved. save_both_split_models = True if len(split_nodes) == 0 else False # split model with given split node split_model_part_1, split_model_part_2 = split_model.split_model_with_node( split_node.name, model.model_path, save_both_split_models, save_path=tmp_file.name ) if not save_both_split_models: # append split_model_part_2 to do next split model_to_split.append(split_model_part_2) logger.info("Quantize split model {}".format(split_idx)) if require_data_reader: # process data_reader for current split and next split current_data_reader = _filter_data_reader_for_current_split_model( split_model_part_1.model, complete_data_reader ) # complete_data_reader contains split_model_part_1 output data complete_data_reader = _prepare_data_reader_for_next_split_model( split_model_part_1.model_path, [i.name for i in split_model_part_2.model.graph.input], complete_data_reader, providers, ) lwq_data_reader.append(complete_data_reader) # perform quantization split_model_part_1_quantized = quant_func( split_model_part_1, weight_config=weight_config, data_reader=current_data_reader, return_modelproto=False, **kwargs ) else: # perform quantization split_model_part_1_quantized = quant_func( split_model_part_1, weight_config=weight_config, return_modelproto=False, **kwargs ) # check split model is valid try: ort.InferenceSession(split_model_part_1_quantized.model_path, providers=providers) except Exception as e: logger.error( "Layer-wise quantized model {} can't be inferred correctly. " "Please check the raise exception".format(split_idx) ) raise e # merge split quantized model if quantized_model_merged is None: quantized_model_merged = split_model_part_1_quantized quantized_model_merged.write_external_data_to_new_location(overwrite=True) else: quantized_model_merged.merge_split_models(split_model_part_1_quantized) split_idx += 1 # if this is the last split, quantize the last split model if save_both_split_models: logger.info("Quantize split model {}".format(split_idx)) # quantize split model if require_data_reader: # process data_reader for current split current_data_reader = lwq_data_reader.pop(0) current_data_reader = _filter_data_reader_for_current_split_model( split_model_part_2.model, complete_data_reader ) # perform quantization split_model_part_2_quantized = quant_func( split_model_part_2, weight_config=weight_config, data_reader=current_data_reader, return_modelproto=False, **kwargs ) else: # perform quantization split_model_part_2_quantized = quant_func( split_model_part_2, weight_config=weight_config, return_modelproto=False, **kwargs ) # check split model is valid try: ort.InferenceSession(split_model_part_2_quantized.model_path, providers=providers) except Exception as e: logger.error( "Layer-wise quantized model {} can't be inferred correctly. " "Please check the raise exception".format(split_idx) ) raise e # merge split quantized model if quantized_model_merged is None: quantized_model_merged = split_model_part_2_quantized quantized_model_merged.write_external_data_to_new_location(overwrite=True) else: quantized_model_merged.merge_split_models(split_model_part_2_quantized) # reload external data to prevent external data file path errors onnx.external_data_helper.load_external_data_for_model( quantized_model_merged.model, os.path.dirname(quantized_model_merged.model_path) ) tmp_file.cleanup() return quantized_model_merged class DataReader(data_reader.CalibrationDataReader): """Data reader for layer-wise quantization.""" def __init__(self, data_list): self.data_list = data_list self.iter_next = iter(self.data_list) def get_next(self): return next(self.iter_next, None) def rewind(self): self.iter_next = iter(self.data_list) def _filter_data_reader_for_current_split_model( model: onnx.ModelProto, current_data_reader: data_reader.CalibrationDataReader, ): """Filter data reader to remove data that is not in model input. Args: model (onnx.ModelProto): onnx model. current_data_reader (data_reader.CalibrationDataReader): data reader of current split model. Returns: data_reader.CalibrationDataReader: filtered data reader. """ filter_inputs = [] input_names = [input.name for input in model.graph.input] current_data_reader.rewind() while True: inputs = current_data_reader.get_next() if not inputs: break filter_input = { input_name: input_tensor for input_name, input_tensor in inputs.items() if input_name in input_names } filter_inputs.append(filter_input) return DataReader(filter_inputs) def _prepare_data_reader_for_next_split_model( model_path: str, next_model_input_names: list, data_reader: data_reader.CalibrationDataReader, providers: List[str] = ["CPUExecutionProvider"], ): """Prepare data reader for next split model. Get data output of current split model and save for next split model. Args: model (str): path to onnx model. data_reader (data_reader.CalibrationDataReader): data reader providers (List[str], optional): providers to use. Defaults to ["CPUExecutionProvider"]. Returns: data_reader.CalibrationDataReader: data reader for next split model. """ data_reader.rewind() data_reader_for_next_split_model = [] session = ort.InferenceSession(model_path, providers=providers) output_names = [output.name for output in session.get_outputs()] input_names = [input.name for input in session.get_inputs()] while True: inputs = data_reader.get_next() if not inputs: break out = session.run(None, {name: inputs[name] for name in input_names}) filter_input = {name: value for name, value in zip(output_names, out)} for name, value in inputs.items(): if name in next_model_input_names and name not in filter_input: filter_input[name] = value data_reader_for_next_split_model.append(filter_input) return DataReader(data_reader_for_next_split_model) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/__init__.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/calibrate.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # ------------------------------------------------------------------------- # Copyright (c) Microsoft, Intel Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- """Calibration for onnx models.""" import copy import logging import os import sys from importlib import util import numpy as np import onnx import onnxruntime from packaging import version from onnx_neural_compressor import logger, onnx_model from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant import calibrator if sys.version_info < (3, 11) and util.find_spec("onnxruntime_extensions"): import onnxruntime_extensions ONNX18_VERSION = version.Version("1.8.0") ORT112_VERSION = version.Version("1.12.0") class ONNXRTAugment: """Augment input model to dump tensor or for calibration.""" def __init__( self, model_wrapper, dataloader, dump_op_types, black_nodes=[], white_nodes=[], iterations=[], execution_provider="CPUExecutionProvider", reduce_range=False, **kwargs, ): """Initialization. Args: model_wrapper (Model): model to be augmented dataloader (object): user implemented object to read in and preprocess calibration dataset dump_op_types (list): operator types to be calibrated and quantized black_nodes (list, optional): operator names that should not be quantized. Defaults to []. white_nodes (list, optional): operator names that force to be quantized. Defaults to []. iterations (list, optional): tensor of which iteration will be collected. Defaults to []. execution_provider (list, optional): execution provider for onnxruntime. Defaults to 'CPUExecutionProvider'. reduce_range (bool, optional): use 7 bit or not. Defaults to False. """ self.model_wrapper = ( model_wrapper if isinstance(model_wrapper, onnx_model.ONNXModel) else onnx_model.ONNXModel(model_wrapper, load_external_data=True) ) self.model = self.model_wrapper.model ai_onnx_domain = [opset for opset in self.model.opset_import if not opset.domain or opset.domain == "ai.onnx"] self.opset_version = ai_onnx_domain[0].version self.dataloader = dataloader self.dump_op_types = dump_op_types self.black_nodes = black_nodes self.white_nodes = white_nodes self.augmented_model = None self.iterations = iterations self.execution_provider = execution_provider self.augment_nodes = [] self.dequantized_output = {} self.already_quantized = "DequantizeLinear" in [node.op_type for node in self.model.graph.node] self.dynamically_quantized = False self.ort_version = version.Version(onnxruntime.__version__) self.reduce_range = reduce_range def augment_graph(self): """Augment_graph. Adds nodes to all quantization_candidates op type nodes in model and ensures their outputs are stored as part of the graph output. Args: activation_only (bool, optional): whether to dump activation tensor only. Defaults to False. weight_only (bool, optional): whether to dump weight_only. Defaults to False. """ self.dequantized_output.clear() onnx_version = version.Version(onnx.__version__) if onnx_version < ONNX18_VERSION: logger.warning("Static quantization for NLP model is supported at onnx 1.8.0 and newer.") if self.already_quantized and any( [i.dims in [1, 2] for i in self.model_wrapper.initializer() if i.name.endswith("_scale")] ): if self.opset_version < 13 and self.ort_version >= ORT112_VERSION: logger.warning( "Please use onnxruntime < 1.12.0 or upgrade model opset " "version to 13 or higher to inspect per-channel quantized weight" ) model = copy.deepcopy(self.model) model_nodes_names = [node.name for node in model.graph.node] added_nodes = [] added_outputs = [] tensors_to_dump = set() for augment_node_type in self.augment_nodes: if augment_node_type not in ["DequantizeLinear"]: # pragma: no cover raise ValueError( "Unexpected augment_node {} only DequantizeLinear is supported".format(augment_node_type) ) if self.already_quantized: # mapping between fp32 node and int8 node new_white_nodes = [] for white_node in self.white_nodes: new_white_node = white_node + "_quant" assert new_white_node in model_nodes_names, "no quantized {} in the graph".format(white_node) new_white_nodes.append(new_white_node) self.white_nodes = new_white_nodes node_outputs = [] for node in model.graph.node: # pylint: disable=no-member node_outputs.extend(node.output) should_be_dump = ((node.op_type in self.dump_op_types) and (node.name not in self.black_nodes)) or ( node.name in self.white_nodes ) if should_be_dump: # add input tensors which should be dump for input in node.input: if len(input) != 0: # to prevent input is "" initializer_tensor = self.model_wrapper.get_initializer(input) if initializer_tensor is None: tensors_to_dump.add(input) # add output tensors which should be dump tensors_to_dump.update([output for output in node.output if len(output) != 0]) model_inputs = [i.name for i in model.graph.input] for tensor in tensors_to_dump: if tensor not in node_outputs and tensor not in model_inputs: continue if self.augment_nodes: for augment_node_type in self.augment_nodes: if augment_node_type in ["DequantizeLinear"]: # insert DequantizeLinear node as output if tensor.endswith("_scale") or tensor.endswith("_zero_point"): # pragma: no cover continue if not self.dynamically_quantized: tensor = ( tensor.replace("_QuantizeInput", "_quantized") if tensor.endswith("_QuantizeInput") else tensor ) else: tensor = ( tensor.replace("_output_quantized", "") if tensor.endswith("_output_quantized") else tensor ) augment_node_name = tensor + "_new_" + augment_node_type scale, zero_point = self.model_wrapper.get_scale_zero(tensor) if scale: # the tensor is in INT8 dtype nodes, output = self._dequantize(tensor, scale, zero_point) if output: added_nodes.extend(nodes) added_outputs.append( onnx.helper.make_tensor_value_info( output, onnx.TensorProto.FLOAT, () # pylint: disable=no-member ) ) # pylint: disable=no-member else: # the tensor is in FP32 dtype if tensor not in [t.name for t in model.graph.output]: added_tensor = onnx.helper.ValueInfoProto() added_tensor.name = tensor added_outputs.append(added_tensor) else: if tensor not in [t.name for t in model.graph.output]: added_tensor = onnx.helper.ValueInfoProto() added_tensor.name = tensor added_outputs.append(added_tensor) if self.augment_nodes: model.graph.node.extend(added_nodes) # pylint: disable=no-member model.graph.output.extend(added_outputs) # pylint: disable=no-member self.augmented_model = model if self.model_wrapper.is_large_model: # pragma: no cover onnx.save_model( model, self.model_wrapper.model_path + "_augment.onnx", save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) def get_activation_tensors_calib_range(self, q_config=None): """Get calib ranges of activation tensors. Args: q_config (dict, optional): quantization config. Defaults to None. Returns: dict: calib ranges """ # conduct inference session and get intermediate outputs so = onnxruntime.SessionOptions() so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL if sys.version_info < (3, 11) and util.find_spec("onnxruntime_extensions"): so.register_custom_ops_library(onnxruntime_extensions.get_library_path()) execution_provider = ( self.execution_provider if self.execution_provider != "TensorrtExecutionProvider" else "CUDAExecutionProvider" ) session = ( onnxruntime.InferenceSession(self.augmented_model.SerializeToString(), so, providers=[execution_provider]) if not self.model_wrapper.is_large_model else onnxruntime.InferenceSession( self.model_wrapper.model_path + "_augment.onnx", so, providers=[execution_provider] ) ) len_inputs = len(session.get_inputs()) inputs_names = [session.get_inputs()[i].name for i in range(len_inputs)] len_outputs = len(session.get_outputs()) outputs_names = [session.get_outputs()[i].name for i in range(len_outputs)] node_output_names = [ output.name if output.name not in self.dequantized_output else self.dequantized_output[output.name] for output in session.get_outputs() ] augment_model_wrapper = ( onnx_model.ONNXModel(self.augmented_model, load_external_data=False) if not self.model_wrapper.is_large_model else onnx_model.ONNXModel(self.model_wrapper.model_path + "_augment.onnx", load_external_data=False) ) input_name_to_nodes = augment_model_wrapper.input_name_to_nodes() output_name_to_node = augment_model_wrapper.output_name_to_node() name_to_node = {} for data_name in node_output_names: node = None if data_name in output_name_to_node: node = output_name_to_node[data_name] elif data_name in input_name_to_nodes: node = input_name_to_nodes[data_name][0] assert node, "{} is neither an input nor an output of nodes in augmented model.".format(data_name) name_to_node[data_name] = node.name activation_tensors_calib_range = {} intermediate_tensor = {} name_to_calibrator = {} ort_inputs_for_next_split_model = [] def _collect_data(inputs): for output_idx, output in enumerate(session.run(None, inputs)): if q_config is not None and output.size != 0: node_name = name_to_node[node_output_names[output_idx]] if node_output_names[output_idx] not in name_to_calibrator: calib_method = ( q_config[node_name]["calibrate_method"] if q_config and node_name in q_config else "MinMax" ) assert calib_method in calibrator.CALIBRATOR, "Calibration method {} is not registered.".format( calib_method ) _calibrator = calibrator.CALIBRATOR[calib_method]() else: _calibrator = name_to_calibrator[node_output_names[output_idx]] # currently, the calibration range for each iteration is collected if # the calibration method is minmax, otherwise the tensor data is collected. # TODO: for entropy and percentile method, need to support range collection # per iteration in the future. if _calibrator.method_name == "MinMax": _calibrator.collect(output) activation_tensors_calib_range[node_output_names[output_idx]] = [list(_calibrator.calib_range)] name_to_calibrator[node_output_names[output_idx]] = _calibrator else: intermediate_tensor.setdefault((node_output_names[output_idx], node_name), []).append(output) elif q_config is None: activation_tensors_calib_range.setdefault(node_output_names[output_idx], []).append(output) idx = 0 while True: inputs = self.dataloader.get_next() if not inputs: break if self.iterations != []: if idx > max(self.iterations): break if idx in self.iterations: _collect_data(inputs) else: _collect_data(inputs) idx += 1 # for entropy and percentile method, collect calibration range after all tensors are collected. merged_dict = intermediate_tensor for (output_name, node_name), datas in merged_dict.items(): if any([data is None for data in datas]): continue if any([data.dtype in [bool] for data in datas]): # output type of some ops is bool, skip continue calib_method = q_config[node_name]["calibrate_method"] if q_config and node_name in q_config else 0 _calibrator = calibrator.CALIBRATOR[calib_method]() _calibrator.collect(datas) activation_tensors_calib_range.setdefault(output_name, []).append(list(_calibrator.calib_range)) _calibrator.clear() del _calibrator return activation_tensors_calib_range def get_weight_tensors_calib_range(self): """Get calib ranges of weight tensors. Returns: dict: calib ranges """ model_nodes_names = [node.name for node in self.model.graph.node] # if augmented_model is not None, it means self.white_nodes is already updated in augment_graph func # then skip update here if self.already_quantized and self.augmented_model is None: # mapping between fp32 node and int8 node new_white_nodes = [] for white_node in self.white_nodes: new_white_node = white_node + "_quant" assert new_white_node in model_nodes_names, "no quantized {} in the " "graph".format(white_node) new_white_nodes.append(new_white_node) self.white_nodes = new_white_nodes added_outputs = set() initializer_tensors_to_dump = [] initializers = [init.name for init in self.model.graph.initializer] for node in self.model.graph.node: # pylint: disable=no-member should_be_dump = ((node.op_type in self.dump_op_types) and (node.name not in self.black_nodes)) or ( node.name in self.white_nodes ) if should_be_dump: for input in node.input: if ( (self.already_quantized and input.replace("_dequantized", "_quantized") in initializers) or (not self.already_quantized and input in initializers) ) and len(input) != 0: added_outputs.add(input) for tensor in added_outputs: if tensor not in initializers: continue if self.augment_nodes: for augment_node_type in self.augment_nodes: if augment_node_type in ["DequantizeLinear"]: if not (tensor.endswith("_scale") or tensor.endswith("_zero_point")): initializer_tensors_to_dump.append(tensor) else: initializer_tensors_to_dump.append(tensor) weight_tensors_calib_range = {} for initializer_tensor_name in initializer_tensors_to_dump: initializer_tensor = self.model_wrapper.get_initializer(initializer_tensor_name) # double check initializer tensor is not None if initializer_tensor is None: # pragma: no cover continue initializer_tensor = onnx.numpy_helper.to_array( initializer_tensor, base_dir=( os.path.dirname(self.model_wrapper.model_path) if self.model_wrapper.model_path is not None else "" ), ) _calibrator = calibrator.CALIBRATOR["MinMax"]() # use minmax method to calibrate initializer tensors if initializer_tensor.flatten().size > 0: _calibrator.collect(initializer_tensor) weight_tensors_calib_range[initializer_tensor_name] = [list(_calibrator.calib_range)] _calibrator.clear() del _calibrator return weight_tensors_calib_range def get_intermediate_outputs(self, q_config=None, activation_only=False, weight_only=False): """Gather intermediate model outputs after running inference.""" output_dicts = {} if not activation_only and not weight_only: output_dicts = self.get_activation_tensors_calib_range(q_config) output_dicts.update(self.get_weight_tensors_calib_range()) elif weight_only: output_dicts = self.get_weight_tensors_calib_range() elif activation_only: output_dicts = self.get_activation_tensors_calib_range(q_config) return list(output_dicts.keys()), output_dicts def _dequantize(self, tensor, scale_tensor, zo_tensor): """Helper function to dequantize tensor.""" int_tensor = self.model_wrapper.get_initializer(tensor) if int_tensor: # weight tensor return self._dequantize_weight(tensor, scale_tensor, zo_tensor) else: return self._dequantize_activation(tensor, scale_tensor, zo_tensor) def _dequantize_activation(self, activation_tensor_name, scale_tensor, zo_tensor): """Helper function to dequantize activation.""" added_nodes, added_output = self._add_dequantize_node(activation_tensor_name, scale_tensor, zo_tensor) self.dequantized_output[added_output] = activation_tensor_name return added_nodes, added_output def _dequantize_weight(self, weight_tensor_name, scale_tensor, zo_tensor): """Helper function to dequantize weight.""" weight_tensor = self.model_wrapper.get_initializer(weight_tensor_name) if len(scale_tensor.dims) in [1, 2] and weight_tensor.dims[0] == max(scale_tensor.dims): logger.debug("weight {} is quantized with per channel granularity.".format(weight_tensor_name)) if self.opset_version < 13 and self.ort_version >= ORT112_VERSION: logger.warning( "Skip dequantizing weight {}, please use onnxruntime < 1.12.0 " "or upgrade model opset version to 13 or higher".format(weight_tensor_name) ) return [], None node = self.model_wrapper.input_name_to_nodes()[weight_tensor_name][0] if "Conv" in node.op_type or ("Gemm" in node.op_type and quant_utils.is_B_transposed(node)): added_nodes, added_output = self._add_dequantize_transpose_node( weight_tensor_name, scale_tensor, zo_tensor, len(weight_tensor.dims) ) else: added_nodes, added_output = self._add_dequantize_node( weight_tensor_name, scale_tensor, zo_tensor, axis=1 if self.opset_version > 12 else None ) else: added_nodes, added_output = self._add_dequantize_node(weight_tensor_name, scale_tensor, zo_tensor) self.dequantized_output[added_output] = weight_tensor_name return added_nodes, added_output def _add_dequantize_node(self, tensor_name, scale_tensor, zo_tensor, axis=None): """Helper function to generate dequantize node.""" dequantize_node = onnx.helper.make_node( "DequantizeLinear", [tensor_name, scale_tensor.name, zo_tensor.name], [tensor_name + "_output"], tensor_name + "_DequantizeLinear", axis, ) return [dequantize_node], tensor_name + "_output" def _add_dequantize_transpose_node(self, tensor_name, scale_tensor, zo_tensor, dim): """Insert Transpose-DequantizelLinear-Transpose pairs.""" pre_transpose_node = onnx.helper.make_node( "Transpose", inputs=[tensor_name], outputs=[tensor_name + "_transposed"], perm=(1, 0, 2, 3) if dim == 4 else (1, 0), name=tensor_name + "_pre_transpose", ) dequantize_node = onnx.helper.make_node( "DequantizeLinear", [tensor_name + "_transposed", scale_tensor.name, zo_tensor.name], [tensor_name + "_DequantizeLinear"], tensor_name + "_DequantizeLinear", axis=1 if self.opset_version > 12 else None, ) post_transpose_node = onnx.helper.make_node( "Transpose", inputs=[tensor_name + "_DequantizeLinear"], outputs=[tensor_name + "_output"], perm=(1, 0, 2, 3) if dim == 4 else (1, 0), name=tensor_name + "_post_transpose", ) added_nodes = [pre_transpose_node, dequantize_node, post_transpose_node] return added_nodes, tensor_name + "_output" def _map_calibration(self, node_output_names, output_dicts): """Map tensor names and min/max values.""" merged_dict = {} for name, minmaxs in output_dicts.items(): for minmax in minmaxs: if len(minmax) < 2: continue merged_dict.setdefault(name + "_Min", []).append(minmax[0]) merged_dict.setdefault(name + "_Max", []).append(minmax[1]) # Characterizing distribution of a node's values across test data sets clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict) pairs = [ tuple([float(min(clean_merged_dict[name + "_Min"])), float(max(clean_merged_dict[name + "_Max"]))]) for name in node_output_names ] final_dict = dict(zip(node_output_names, pairs)) return final_dict def dump_minmax(self, q_config): """Get calib ranges of tensors.""" # pipeline of getting calib ranges of tensors during calibration: # 1. augment_graph(): insert activation tensors to model output # 2. get_intermediate_outputs(): # 2.1 get_activation_tensors_calib_range(): get calib ranges of activation tensors using the augment graph # 2.2 get_weight_tensors_calib_range(): get calib ranges of weight tensors self.augment_graph() node_output_names, output_dicts = self.get_intermediate_outputs(q_config) return self._map_calibration(node_output_names, output_dicts) def dump_calibration(self, q_config, min_max=None): """Gather calibration params for quantization. Args: q_config (dict): op-wise quantization config min_max (dict, optional): min/max values of tensors """ return ( self.calculate_quantization_params(q_config, self.dump_minmax(q_config)) if min_max is None else self.calculate_quantization_params(q_config, min_max) ) def calculate_quantization_params(self, q_config, quantization_thresholds): """Given quantization thresholds, calculate the quantization params. Args: q_config (dict): op-wise quantization config quantization_thresholds (dict): Dictionary specifying the min and max values or outputs of conv and matmul nodes, should be specified in the following format: {"param_name": [min, max]} """ if quantization_thresholds is None: raise ValueError( "quantization thresholds is required to calculate quantization \ params (zero point and scale)" ) quantization_params = {} model = self.model input_name_to_nodes = self.model_wrapper.input_name_to_nodes() output_name_to_node = self.model_wrapper.output_name_to_node() for tensor_name in quantization_thresholds.keys(): child = None if tensor_name in input_name_to_nodes: children = input_name_to_nodes[tensor_name] if len(children) == 1: child = children[0] parent = None sym = False qType = 2 # uint8 # input and output tensor follow activation_type and activation_sym if tensor_name in input_name_to_nodes and any( [i.name in q_config for i in input_name_to_nodes[tensor_name]] ): for child in input_name_to_nodes[tensor_name]: if child.name in q_config and q_config[child.name] not in ["fp32", "fp16", "bf16"]: sym = q_config[child.name]["activation_sym"] qType = q_config[child.name]["activation_type"] break elif ( tensor_name in output_name_to_node and output_name_to_node[tensor_name].name in q_config and q_config[output_name_to_node[tensor_name].name] not in ["fp32", "fp16", "bf16"] ): sym = q_config[output_name_to_node[tensor_name].name]["activation_sym"] qType = q_config[output_name_to_node[tensor_name].name]["activation_type"] if self.execution_provider in ["TensorrtExecutionProvider"]: # TensorrtExecutionProvider only support int8 qType = 3 node_thresholds = quantization_thresholds[tensor_name] node_params = self.calculate_scale_zeropoint( parent, child, node_thresholds[0], node_thresholds[1], sym, qType, ) quantization_params[tensor_name] = node_params return quantization_params def calculate_scale_zeropoint(self, last_node, next_node, rmin, rmax, sym, qType): """Given the source and destination node of tensor, return calculated zero point and scales.""" zp_and_scale = [] # adjust rmin and rmax such that 0 is included in the range. This is required # to make sure zero can be uniquely represented. rmin = min(rmin, 0) rmax = max(rmax, 0) if next_node: if next_node.op_type == "Relu": if rmin < 0: rmin = 0 elif next_node.op_type == "Clip" and len(next_node.input) == 3: if self.model_wrapper.get_initializer(next_node.input[1]) is not None: clip_min = onnx.numpy_helper.to_array(self.model_wrapper.get_initializer(next_node.input[1])) if rmin < clip_min: rmin = clip_min.tolist() if not isinstance(clip_min.tolist(), list) else clip_min.tolist()[0] if self.model_wrapper.get_initializer(next_node.input[2]) is not None: clip_max = onnx.numpy_helper.to_array(self.model_wrapper.get_initializer(next_node.input[2])) if rmax > clip_max: rmax = clip_max.tolist() if not isinstance(clip_max.tolist(), list) else clip_max.tolist()[0] if last_node: if last_node.op_type in ["Conv", "FusedConv"]: attrs = [attr for attr in last_node.attribute] attrs_names = [attr.name for attr in last_node.attribute] if "activation" in attrs_names: if attrs[attrs_names.index("activation")].s == b"Relu": rmin = max(rmin, 0) if attrs[attrs_names.index("activation")].s == b"Clip": assert ( "activation_params" in attrs_names ), "the model contains no params for clip node {}".format(last_node) clip_params = attrs[attrs_names.index("activation_params")].floats rmin = min(rmin, clip_params[0], clip_params[1]) rmax = max(rmax, clip_params[0], clip_params[1]) scale, zp = quant_utils.calculate_scale_zp(rmin, rmax, qType, sym, self.reduce_range) zp_and_scale.append(zp) zp_and_scale.append(scale) return zp_and_scale ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/calibrator.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # ------------------------------------------------------------------------- # Copyright (c) Microsoft, Intel Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- """Calibrator for onnx models.""" import copy import numpy as np from scipy import stats CALIBRATOR = {} def calib_registry(calib_method): """The class decorator used to register all Calibrator subclasses.""" def decorator_calib(cls): assert cls.__name__.endswith( "Calibrator" ), "The name of subclass of Calibrator should end with 'Calibrator' substring." if cls.__name__[: -len("Calibrator")] in CALIBRATOR: # pragma: no cover raise ValueError("Cannot have two operators with the same name.") CALIBRATOR[calib_method] = cls return cls return decorator_calib class CalibratorBase: """Base calibrator class.""" def __init__(self): """Initialize base calibrator class.""" self._calib_min = None self._calib_max = None def collect(self, datas): """Collect calibration range.""" self.collect_calib_data(datas) def clear(self): """Clear calibration range.""" self._calib_min = None self._calib_max = None def collect_calib_data(self, datas): """Collect calibration range value.""" raise NotImplementedError @property def calib_range(self): """Get calibration range value.""" return self._calib_min, self._calib_max @calib_registry(calib_method="MinMax") class MinMaxCalibrator(CalibratorBase): """MinMax calibrator class.""" def __init__(self): """Initialize minmax calibrator class.""" super(MinMaxCalibrator, self).__init__() def collect_calib_data(self, datas): """Collect calibration range.""" if isinstance(datas, list) and len(set([data.shape for data in datas])) != 1: for data in datas: if data.size == 0: # pragma: no cover continue self._collect_value(data) else: datas = np.asarray(datas) datas = datas.flatten() assert datas.size > 0, "collected intermediate data size" "should not be 0, please check augmented_model" self._collect_value(datas) def _collect_value(self, data): """Collect min/max value.""" data = np.asarray(data) local_min = np.min(data[np.isinf(data) == False]) # noqa: E712 local_max = np.max(data[np.isinf(data) == False]) # noqa: E712 if self._calib_min is None and self._calib_max is None: self._calib_min = local_min self._calib_max = local_max else: self._calib_min = np.minimum(self._calib_min, local_min) self._calib_max = np.maximum(self._calib_max, local_max) @property def method_name(self): """Get calibration method name.""" return "MinMax" @calib_registry(calib_method="Percentile") class PercentileCalibrator(CalibratorBase): """Percentile calibrator class. Args: num_bins (int, optional): number of bins to create a new histogram for collecting tensor values. Defaults to 2048. percentile (float, optional): A float number between [0, 100]. Defaults to 99.999. """ def __init__(self, num_bins=2048, percentile=99.999): """Initialize percentile calibrator class.""" super(PercentileCalibrator, self).__init__() self.collector = None self.num_bins = num_bins self.percentile = percentile def collect_calib_data(self, datas): """Collect calibration range.""" if not self.collector: self.collector = HistogramCollector(self.num_bins) self.collector.collect_data(datas) self.compute_percentile_range(self.percentile) def compute_percentile_range(self, percentile): """Compute percentile range.""" if percentile < 0 or percentile > 100: raise ValueError("Invalid percentile. Must be in range 0 <= percentile <= 100.") calib_hist, calib_bin_edges, min_range, max_range, th = self.collector.histogram total = calib_hist.sum() cdf = np.cumsum(calib_hist / total) percent_to_cut_one_side = (100.0 - percentile) / 200.0 max_idx = np.searchsorted(cdf, 1.0 - percent_to_cut_one_side) min_idx = np.searchsorted(cdf, percent_to_cut_one_side) self._calib_min = calib_bin_edges[min_idx].astype("float32") self._calib_max = calib_bin_edges[max_idx].astype("float32") if self._calib_min < min_range: self._calib_min = min_range if self._calib_max > max_range: self._calib_max = max_range def clear(self): """Clear calibration range.""" self._calib_min = None self._calib_max = None self.collector = None @property def method_name(self): """Get calibration method name.""" return "Percentile" @calib_registry(calib_method="Entropy") class EntropyCalibrator(CalibratorBase): """Entropy calibrator class. Args: num_bins (int, optional):number of bins to create a new histogram for collecting tensor values. Defaults to 128. num_quantized_bins (int, optional): number of quantized bins. Defaults to 128. """ def __init__(self, num_bins=128, num_quantized_bins=128): """Initialize entropy calibrator class.""" super(EntropyCalibrator, self).__init__() self.collector = None self.num_bins = num_bins self.num_quantized_bins = num_quantized_bins def collect_calib_data(self, datas): """Collect calibration range.""" if not self.collector: self.collector = HistogramCollector(self.num_bins) self.collector.collect_data(datas) self.compute_kl_range() def compute_kl_range(self): """Compute entropy range.""" histogram = self.collector.histogram self._calib_min, self._calib_max = self.get_kl_threshold(histogram, self.num_quantized_bins) def get_kl_threshold(self, histogram, num_quantized_bins): """Compute entropy threshold. Ref: https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/quantization/calibrate.py Args: histogram (tuple): hist, hist_edges, min, max and threshold num_quantized_bins (int): number of quantized bins. Returns: float: optimal threshold """ hist = histogram[0] hist_edges = histogram[1] num_bins = hist.size zero_bin_index = num_bins // 2 num_half_quantized_bin = num_quantized_bins // 2 kl_divergence = np.zeros(zero_bin_index - num_half_quantized_bin + 1) thresholds = [(0, 0) for i in range(kl_divergence.size)] for i in range(num_half_quantized_bin, zero_bin_index + 1, 1): start_index = zero_bin_index - i end_index = zero_bin_index + i + 1 if (zero_bin_index + i + 1) <= num_bins else num_bins thresholds[i - num_half_quantized_bin] = ( float(hist_edges[start_index]), float(hist_edges[end_index]), ) sliced_distribution = copy.deepcopy(hist[start_index:end_index]) # reference distribution p p = sliced_distribution.copy() # a copy of np array left_outliers_count = sum(hist[:start_index]) right_outliers_count = sum(hist[end_index:]) p[0] += left_outliers_count p[-1] += right_outliers_count # nonzeros[i] incidates whether p[i] is non-zero nonzeros = (p != 0).astype(np.int64) # quantize p.size bins into quantized bins (default 128 bins) quantized_bins = np.zeros(num_quantized_bins, dtype=np.int64) num_merged_bins = sliced_distribution.size // num_quantized_bins # merge bins into quantized bins for index in range(num_quantized_bins): start = index * num_merged_bins end = start + num_merged_bins quantized_bins[index] = sum(sliced_distribution[start:end]) quantized_bins[-1] += sum(sliced_distribution[num_quantized_bins * num_merged_bins :]) # in order to compare p and q, we need to make length of q equals to length of p # expand quantized bins into p.size bins q = np.zeros(p.size, dtype=np.int64) for index in range(num_quantized_bins): start = index * num_merged_bins end = start + num_merged_bins norm = sum(nonzeros[start:end]) if norm != 0: q[start:end] = float(quantized_bins[index]) / float(norm) p = smooth_distribution(p) q = smooth_distribution(q) if isinstance(q, np.ndarray): kl_divergence[i - num_half_quantized_bin] = stats.entropy(p, q) else: kl_divergence[i - num_half_quantized_bin] = float("inf") min_kl_divergence_idx = np.argmin(kl_divergence) optimal_threshold = thresholds[min_kl_divergence_idx] min_value = histogram[2] max_value = histogram[3] if optimal_threshold[0] < min_value: optimal_threshold = (min_value, optimal_threshold[1]) if optimal_threshold[1] > max_value: optimal_threshold = (optimal_threshold[0], max_value) return optimal_threshold[0], optimal_threshold[1] def clear(self): """Clear calibration range.""" self._calib_min = None self._calib_max = None self.collector = None @property def method_name(self): """Get calibration method name.""" return "Entropy" class HistogramCollector: """Histogram collctor class.""" def __init__(self, num_bins=2048): """Initialize histogram collctor.""" self._num_bins = num_bins self._histogram = None def collect_data(self, datas): """Collect histogram data.""" if isinstance(datas, list) and len(set([data.shape for data in datas])) != 1: for data in datas: if data.size == 0: # pragma: no cover continue self._collect_value(data) else: datas = np.asarray(datas) datas = datas.flatten() assert datas.size > 0, "collected intermediate data size" "should not be 0, please check augmented_model" self._collect_value(datas) def _collect_value(self, data): """Collect value.""" data = np.asarray(data) min_range = np.min(data) max_range = np.max(data) th = max(abs(min_range), abs(max_range)) if self._histogram is None: hist, hist_edges = np.histogram(data, self._num_bins, range=(-th, th)) self._histogram = (hist, hist_edges, min_range, max_range, th) else: self._histogram = self.combine_histogram(self._histogram, data, min_range, max_range, th) def combine_histogram(self, old_hist, data_arr, new_min, new_max, new_th): """Combine histogram.""" (old_hist, old_hist_edges, old_min, old_max, old_th) = old_hist if new_th <= old_th: hist, _ = np.histogram(data_arr, bins=len(old_hist), range=(-old_th, old_th)) return ( old_hist + hist, old_hist_edges, min(old_min, new_min), max(old_max, new_max), old_th, ) else: # Need to generate new histogram with new_th if old_th == 0: hist, hist_edges = np.histogram(data_arr, len(old_hist), range=(-new_th, new_th)) hist += old_hist else: old_num_bins = len(old_hist) old_step = 2 * old_th / old_num_bins half_increased_bins = int((new_th - old_th) // old_step + 1) new_num_bins = half_increased_bins * 2 + old_num_bins new_th = half_increased_bins * old_step + old_th hist, hist_edges = np.histogram(data_arr, bins=new_num_bins, range=(-new_th, new_th)) hist[half_increased_bins : new_num_bins - half_increased_bins] += old_hist return ( hist, hist_edges, min(old_min, new_min), max(old_max, new_max), new_th, ) @property def histogram(self): """Get histogram.""" return self._histogram def smooth_distribution(p, eps=0.0001): """Smooth distribution. Given a discrete distribution (may have not been normalized to 1), smooth it by replacing zeros with eps multiplied by a scaling factor and taking the corresponding amount off the non-zero values. Ref: http://hanj.cs.illinois.edu/cs412/bk3/KL-divergence.pdf https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/quantization/calibrate.py Args: p (array): distribution array eps (float, optional): a small probability. Defaults to 0.0001. Returns: array: smoothed distribution """ is_zeros = (p == 0).astype(np.float32) is_nonzeros = (p != 0).astype(np.float32) n_zeros = is_zeros.sum() n_nonzeros = p.size - n_zeros if not n_nonzeros: return -1 eps1 = eps * float(n_zeros) / float(n_nonzeros) assert eps1 < 1.0, "n_zeros=%d, n_nonzeros=%d, eps1=%f" % ( n_zeros, n_nonzeros, eps1, ) hist = p.astype(np.float32) hist += eps * is_zeros + (-eps1) * is_nonzeros assert (hist <= 0).sum() == 0 return hist ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/__init__.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Operators for onnx model.""" import glob from os import path from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op modules = glob.glob(path.join(path.dirname(__file__), "*.py")) for f in modules: if path.isfile(f) and not f.startswith("__") and not f.endswith("__init__.py"): __import__(path.basename(f)[:-3], globals(), locals(), level=1) OPERATORS = base_op.OPERATORS ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/activation.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Activation operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="LeakyRelu, Sigmoid", mode=[constants.STATIC_QUANT]) class ActivationOperator(base_op.Operator): """Activation operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ActivationOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node data_found, _, _, _, _ = self.quantizer._get_quantization_params(node.output[0]) if not data_found: return False return True def quantize(self): """Do quantizaion.""" node = self.node super().quantize() node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node children = self.quantizer.model.get_children(node) if len(children) == 0 or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parent = self.quantizer.model.get_parents(node)[0] child = self.quantizer.model.get_children(node)[0] inputs = [] inputs.extend(parent.input) inputs.extend(child.input[1:]) qlinear_activation_output = child.output[0] kwargs = {} for attribute in node.attribute: # pragma: no cover kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qlinear_activation_node = onnx.helper.make_node( "QLinear" + node.op_type, inputs, [qlinear_activation_output], node.name, **kwargs ) self.quantizer.new_nodes.append(qlinear_activation_node) self.quantizer.remove_nodes.extend([parent, child, node]) @base_op.op_registry(op_types="Relu, Clip", mode=[constants.STATIC_QUANT]) class RemovableActivationOperator(base_op.Operator): """Removable activation operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(RemovableActivationOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if node.input[0] not in self.quantizer.quantized_value_map: return False return True def quantize(self): """Do quantization.""" node = self.node if node.output[0] in [i.name for i in self.quantizer.model.model.graph.output]: self.quantizer.dequantize_tensor(node, node.input[0]) else: self.quantizer.model.replace_input_of_all_nodes(node.output[0], node.input[0]) self.quantizer.remove_nodes.append(node) @base_op.op_registry( op_types="Softmax, BiasGelu, Elu, Exp, FastGelu, Gelu, Softplus, Tanh", mode=[constants.STATIC_QUANT] ) class Float16ActivationOperator(base_op.Operator): """Float16 Activation operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(Float16ActivationOperator, self).__init__(onnx_quantizer, onnx_node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/argmax.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ArgMax operator.""" from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="ArgMax", mode=[constants.STATIC_QUANT]) class ArgMaxOperator(base_op.Operator): """ArgMax operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ArgMaxOperator, self).__init__(onnx_quantizer, onnx_node) def convert_check(self): """Check if conversion can be done.""" node = self.node return True def convert(self): """Convert to quantized format.""" node = self.node origin_name = node.input[0].split("_argmax_node")[0] if origin_name in self.quantizer.quantized_value_map: node.name = node.name + "_quant" ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/attention.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Attention operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Attention", mode=[constants.DYNAMIC_QUANT, constants.STATIC_QUANT]) class AttentionOperator(base_op.Operator): """Attention operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(AttentionOperator, self).__init__(onnx_quantizer, onnx_node) def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0, 1]) node.name = node.name + "_quant" def convert(self): """Convert QDQ mode to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) quantized_name = [] scale = [] zp = [] for parent in parents[:2]: if parent.op_type == "DynamicQuantizeLinear": quantized_name.append(parent.output[0]) scale.append(parent.output[1]) zp.append(parent.output[2]) elif parent.op_type == "DequantizeLinear": quantized_name.append(parent.input[0]) scale.append(parent.input[1]) zp.append(parent.input[2]) self.quantizer.remove_nodes.append(parent) inputs = [] inputs.extend(quantized_name) inputs.append(node.input[2]) inputs.extend(scale) inputs.append(node.input[3] if len(node.input) > 3 else "") inputs.extend(zp) if len(node.input) > 4: inputs.append(node.input[4]) kwargs = {} for attribute in node.attribute: # pragma: no cover kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qattention_node = onnx.helper.make_node("QAttention", inputs, node.output, node.name, **kwargs) self.quantizer.new_nodes.append(qattention_node) self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/base_op.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Base Operator.""" from onnx_neural_compressor import constants, quantization OPERATORS = { "dynamic_quant": {}, "static_quant": {}, } def op_registry(op_types, mode): """The class decorator used to register all Operator subclasses.""" def decorator_op(cls): assert cls.__name__.endswith( "Operator" ), "The name of subclass of Operator should end with 'Operator' substring." for item in mode: if cls.__name__[: -len("Operator")] in OPERATORS[item]: # pragma: no cover raise ValueError("Cannot have two operators with the same name for {} mode.".format(item)) break for single_op_type in [op_type.strip() for op_type in op_types.split(",")]: for item in mode: OPERATORS[item][single_op_type] = cls return cls return decorator_op class Operator(object): """Base Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" self.quantizer = onnx_quantizer self.node = onnx_node node_name = self.node.name.split("_quant")[0] if node_name in self.quantizer.config: self.dtype = self.quantizer.config[node_name] self.disable_qdq_for_node_output = ( True if onnx_node.op_type in onnx_quantizer.optypes_to_exclude_output_quant else False ) self.per_channel = False self.calibrate_method = 0 # minmax self.weight_sym = True self.weight_dtype = None self.activation_dtype = None self.activation_sym = False if node_name in self.quantizer.config: if self.quantizer.config[node_name] not in self.quantizer.fallback_list: self.per_channel = self.quantizer.config[node_name]["per_channel"] self.calibrate_method = self.quantizer.config[node_name]["calibrate_method"] self.weight_sym = self.quantizer.config[node_name]["weight_sym"] self.weight_dtype = self.quantizer.config[node_name]["weight_type"] self.activation_dtype = self.quantizer.config[node_name]["activation_type"] self.activation_sym = self.quantizer.config[node_name]["activation_sym"] def quantize_check(self): """Check if quantizaion can be done.""" return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node) if not self.disable_qdq_for_node_output or self.quantizer.mode != constants.DYNAMIC_QUANT: self.quantizer.quantize_outputs(node) def convert_check(self): """Check if conversion can be done.""" node = self.node if not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" return ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/binary_op.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Binary operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Add, Mul", mode=[constants.STATIC_QUANT]) class BinaryOperator(base_op.Operator): """Binary operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(BinaryOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node data_found, _, _, _, _ = self.quantizer._get_quantization_params(node.output[0]) if not data_found: return False if self.quantizer.execution_provider == "TensorrtExecutionProvider": return True if not all([self.quantizer.is_valid_quantize_weight(i) for i in node.input]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, initializer_use_weight_qType=False) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node children = self.quantizer.model.get_children(node) if len(children) == 0 or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) child = self.quantizer.model.get_children(node)[0] qlinear_binary_math_output = child.output[0] kwargs = {} for attribute in node.attribute: # pragma: no cover kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qlinear_binary_math_inputs = [] for parent in parents: qlinear_binary_math_inputs.extend(parent.input) qlinear_binary_math_inputs.extend(child.input[1:]) qlinear_binary_math_node = onnx.helper.make_node( "QLinear" + node.op_type, qlinear_binary_math_inputs, [qlinear_binary_math_output], node.name, **kwargs ) self.quantizer.new_nodes += [qlinear_binary_math_node] self.quantizer.remove_nodes.extend(parents) self.quantizer.remove_nodes.append(child) self.quantizer.remove_nodes.append(node) @base_op.op_registry(op_types="Mod", mode=[constants.STATIC_QUANT]) class BinaryDirect8BitOperator(base_op.Operator): """Binary operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(BinaryDirect8BitOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node data_found, _, _, _, _ = self.quantizer._get_quantization_params(node.output[0]) if not data_found: return False if not all([self.quantizer.is_valid_quantize_weight(i) for i in node.input]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, initializer_use_weight_qType=False) if not self.disable_qdq_for_node_output or self.quantizer.mode != "qdq": self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node children = self.quantizer.model.get_children(node) if len(children) == 0 or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if any([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): for idx, parent in enumerate(parents): if parent.op_type == "DequantizeLinear": self.node.input[idx] = parent.input[0] self.quantizer.remove_nodes.append(parent) for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], node.output[0] + "_quantized") node.output[0] = node.output[0] + "_quantized" @base_op.op_registry( op_types="Sum, Sub, Div, Pow, Equal, Greater, GreaterOrEqual, Less, LessOrEqual", mode=[constants.STATIC_QUANT] ) class Float16BinaryOperator(base_op.Operator): """Float16 Binary operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(Float16BinaryOperator, self).__init__(onnx_quantizer, onnx_node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/concat.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Concat Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Concat", mode=[constants.STATIC_QUANT]) class ConcatOperator(base_op.Operator): """Concat Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ConcatOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if len(node.input) == 1: # pragma: no cover return False inits = [i.name for i in self.quantizer.model.initializer()] if all([inp not in self.quantizer.quantized_value_map and inp not in inits for inp in node.input]) or not all( [inp in self.quantizer.quantized_value_map or inp in inits for inp in node.input] ): return False return True def quantize(self): """Do quantizaion.""" node = self.node inits = [i.name for i in self.quantizer.model.initializer()] for idx, inp in enumerate(node.input): initializer_use_weight_qType = inp not in inits self.quantizer.quantize_inputs(node, [idx], initializer_use_weight_qType) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if len(children) == 0 or len(parents) == 0 or not node.name.endswith("_quant"): return False # check input type if all([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): input_zp, input_scale, output_zp = [], [], [] input_zp = [parent.input[2] for parent in parents] input_scale = [parent.input[1] for parent in parents] output_zp = [child.input[2] for child in children if child.op_type == "QuantizeLinear"] if ( any([self.quantizer.model.get_initializer(zp) is None for zp in input_zp]) or any([self.quantizer.model.get_initializer(zp) is None for zp in output_zp]) or any([self.quantizer.model.get_initializer(scale) is None for scale in input_scale]) ): # pragma: no cover return False # check input scale is float type if any( [self.quantizer.model.get_initializer(scale).data_type != 1 for scale in input_scale] ): # pragma: no cover return False # check input zp type is the same with output zp type if any( [ self.quantizer.model.get_initializer(in_zp).data_type not in [self.quantizer.model.get_initializer(out_zp).data_type for out_zp in output_zp] for in_zp in input_zp ] ): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if all([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): inputs = [] inputs.extend([i for i in children if i.op_type == "QuantizeLinear"][0].input[1:]) for parent in parents: inputs.extend(parent.input) self.quantizer.remove_nodes.append(parent) for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], node.output[0] + "_quantized") kwargs = {} for attribute in node.attribute: kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qlconcat_node = onnx.helper.make_node( "QLinearConcat", inputs, [node.output[0] + "_quantized"], node.name, **kwargs ) self.quantizer.new_nodes += [qlconcat_node] self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/conv.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Conv Operator.""" import onnx from onnx import onnx_pb as onnx_proto from onnx_neural_compressor import constants from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Conv, FusedConv", mode=[constants.DYNAMIC_QUANT]) class ConvOperator(base_op.Operator): """Conv Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ConvOperator, self).__init__(onnx_quantizer, onnx_node) def quantize(self): """Do quantizaion.""" node = self.node if node.op_type == "FusedConv": kwargs = {} for attribute in node.attribute: if attribute.name == "activation" and attribute.s in [b"Relu", b"Clip"]: continue if attribute.name == "activation_params": continue kwargs.update(quant_utils.attribute_to_kwarg(attribute)) conv = onnx.helper.make_node("Conv", node.input, node.output, node.name, **kwargs) node.CopyFrom(conv) self.quantizer.quantize_inputs(node, [0]) if self.per_channel: self.quantizer.quantize_weights_per_channel(node, [1], self.weight_dtype, self.weight_sym, 0) else: self.quantizer.quantize_inputs(node, [1]) if len(node.input) == 3: self.quantizer.quantize_bias_tensor(node) node.name = node.name + "_quant" def convert(self): """Convert to QOperator format.""" node = self.node inputs = [] parents = self.quantizer.model.get_parents(node) if parents[0].op_type == "QuantizeLinear": inputs.append(parents[0].output[0]) inputs.append(parents[1].input[0]) inputs.append(parents[0].input[2]) inputs.append(parents[1].input[2]) scale_0 = parents[0].input[1] else: inputs.append(parents[0].output[0]) inputs.append(parents[1].input[0]) inputs.append(parents[0].output[2]) inputs.append(parents[1].input[2]) scale_0 = parents[0].output[1] scale_1 = parents[1].input[1] # quantize bias if exist quantized_bias_name = "" bias_present = False if len(node.input) == 3: quantized_bias_name = node.input[2] + "_quantized" bias_present = True conv_integer_output = node.output[0] + "_output_quantized" kwargs = {} for attribute in node.attribute: if attribute.name == "activation" and attribute.s in [b"Relu", b"Clip"]: # pragma: no cover continue if attribute.name == "activation_params": # pragma: no cover continue kwargs.update(quant_utils.attribute_to_kwarg(attribute)) conv_integer_node = onnx.helper.make_node("ConvInteger", inputs, [conv_integer_output], node.name, **kwargs) self.quantizer.new_nodes.append(conv_integer_node) # Add bias add nodes if bias_present: conv_integer_output = self.quantizer.get_bias_add_nodes( node, parents[1].input[0], conv_integer_output, quantized_bias_name ) # Add cast operation to cast convInteger output to float. cast_op_output = conv_integer_output + "_cast_output" cast_node = onnx.helper.make_node( "Cast", [conv_integer_output], [cast_op_output], conv_integer_output + "_cast", to=onnx_proto.TensorProto.FLOAT, ) self.quantizer.new_nodes.append(cast_node) # Add mul operation to multiply scales of two inputs. scales_mul_op = node.name + "_scales_mul" scales_mul_node = quant_utils.find_by_name(scales_mul_op, self.quantizer.new_nodes) if scales_mul_node is None: scales_mul_node = onnx.helper.make_node("Mul", [scale_0, scale_1], [scales_mul_op + ":0"], scales_mul_op) self.quantizer.new_nodes.append(scales_mul_node) scales_mul_op_output = scales_mul_node.output[0] # Add mul operation to multiply mul_scales_op result with output of ConvInteger # and make the output of this node the same as output of original conv node. output_scale_mul_op = node.name + "_output_scale_mul" self.quantizer.new_nodes.append( onnx.helper.make_node("Mul", [cast_op_output, scales_mul_op_output], [node.output[0]], output_scale_mul_op) ) self.quantizer.remove_nodes.extend(parents[1:]) self.quantizer.remove_nodes.append(node) @base_op.op_registry(op_types="Conv, FusedConv", mode=[constants.STATIC_QUANT]) class StaticConvOperator(ConvOperator): """Conv Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ConvOperator, self).__init__(onnx_quantizer, onnx_node) def quantize(self): """Do quantizaion.""" node = self.node if node.op_type == "FusedConv": kwargs = {} for attribute in node.attribute: if attribute.name == "activation" and attribute.s in [b"Relu", b"Clip"]: continue if attribute.name == "activation_params": continue kwargs.update(quant_utils.attribute_to_kwarg(attribute)) conv = onnx.helper.make_node("Conv", node.input, node.output, node.name, **kwargs) node.CopyFrom(conv) self.quantizer.quantize_inputs(node, [0]) if self.per_channel: self.quantizer.quantize_weights_per_channel(node, [1], self.weight_dtype, self.weight_sym, 0) else: self.quantizer.quantize_inputs(node, [1]) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(node) if len(node.input) == 3: self.quantizer.quantize_bias_tensor(node) node.name = node.name + "_quant" def convert(self): """Convert to QOperator format.""" node = self.node if len(self.quantizer.model.get_children(node)) == 0 or not node.name.endswith("_quant"): # pragma: no cover return parents = self.quantizer.model.get_parents(node) child = self.quantizer.model.get_children(node)[0] qlinear_conv_inputs = [] for parent in parents[0:2]: qlinear_conv_inputs.extend(parent.input) qlinear_conv_inputs.extend(child.input[1:]) if len(parents) == 3: qlinear_conv_inputs.append(parents[-1].input[0]) qlinear_conv_output = child.output[0] kwargs = {} for attribute in node.attribute: if attribute.name == "activation" and attribute.s in [b"Relu", b"Clip"]: # pragma: no cover continue if attribute.name == "activation_params": # pragma: no cover continue kwargs.update(quant_utils.attribute_to_kwarg(attribute)) qlinear_conv_node = onnx.helper.make_node( "QLinearConv", qlinear_conv_inputs, [qlinear_conv_output], node.name, **kwargs ) self.quantizer.new_nodes.append(qlinear_conv_node) self.quantizer.remove_nodes.extend(parents) self.quantizer.remove_nodes.append(child) self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/direct_q8.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Direct8Bit Operator.""" from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry( op_types="Reshape, Transpose, Squeeze, Unsqueeze, Flatten, Expand, Slice, " "SpaceToDepth, DepthToSpace, Upsample, Tile, CenterCropPad", mode=[constants.STATIC_QUANT], ) class Direct8BitOperator(base_op.Operator): """Direct8Bit Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(Direct8BitOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if not self.quantizer.is_valid_quantize_weight(node.input[0]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(self.node, [0], initializer_use_weight_qType=False, direct_int8=True) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(self.node, direct_int8=True) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if (len(children) == 0 and len(parents) == 0) or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if any([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): for parent in parents: if parent.op_type == "DequantizeLinear": # make sure parent DequantizeLinear of input 0 is not used by other ops if len(self.quantizer.model.get_children(parent)) == 1 and not self.quantizer.model.is_graph_output( parents[0].output[0] ): self.quantizer.remove_nodes.append(parent) self.node.input[0] = parent.input[0] break for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], node.output[0] + "_quantized") node.output[0] = node.output[0] + "_quantized" ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/embed_layernorm.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """EmbedLayerNormalization Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="EmbedLayerNormalization", mode=[constants.DYNAMIC_QUANT, constants.STATIC_QUANT]) class EmbedLayerNormalizationOperator(base_op.Operator): """EmbedLayerNormalization Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(EmbedLayerNormalizationOperator, self).__init__(onnx_quantizer, onnx_node) def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [2, 3, 4, 5, 6]) node.name = node.name + "_quant" def convert(self): """Convert to QOperator format.""" node = self.node parents = [i for i in self.quantizer.model.get_parents(node) if i.op_type == "DequantizeLinear"] inputs = [] # 'input_ids' inputs.extend([node.input[0]]) # 'segment_ids' inputs.extend([node.input[1]]) for parent in parents: inputs.append(parent.input[0]) # 'mask' (optional) if len(node.input) > 7: inputs.append(node.input[7]) for parent in parents: inputs.append(parent.input[1]) for parent in parents: inputs.append(parent.input[2]) kwargs = {} for attribute in node.attribute: # pragma: no cover kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qembed_layer_norm_node = onnx.helper.make_node( "QEmbedLayerNormalization", inputs, node.output, node.name, **kwargs ) self.quantizer.new_nodes.append(qembed_layer_norm_node) self.quantizer.remove_nodes.extend(parents) self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/gather.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Gather Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry( op_types="Gather, GatherElements, GatherND", mode=[constants.DYNAMIC_QUANT, constants.STATIC_QUANT] ) class GatherOperator(base_op.Operator): """Gather Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(GatherOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if not self.quantizer.is_valid_quantize_weight(node.input[0]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0], initializer_use_weight_qType=False) if not self.disable_qdq_for_node_output or self.quantizer.mode != constants.DYNAMIC_QUANT: self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if len(children) == 0 or len(parents) == 0 or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" # DQ-Gather-Q-DQ-op node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if any([i.op_type == "DequantizeLinear" for i in parents]): inputs = [] inputs.append(parents[0].input[0]) inputs.append(node.input[1]) out_scale = 1.0 out_zp = 0 gather_new_output = node.output[0] + "_quantized" # dynamic quant output name for child in children: if child.op_type == "QuantizeLinear": out_scale = onnx.numpy_helper.to_array(self.quantizer.model.get_initializer(children[0].input[1])) out_zp = onnx.numpy_helper.to_array(self.quantizer.model.get_initializer(children[0].input[2])) gather_new_output = children[0].output[0] # static quant output name self.quantizer.remove_nodes.append(child) kwargs = {} for attribute in node.attribute: # pragma: no cover kwargs.update(quant_utils.attribute_to_kwarg(attribute)) gather_node = onnx.helper.make_node(node.op_type, inputs, [gather_new_output], node.name, **kwargs) self.quantizer.new_nodes.append(gather_node) if any([i.op_type != "QuantizeLinear" for i in children]): dq_inputs = [] dq_inputs.append(gather_new_output) dq_inputs.extend(parents[0].input[1:]) dq_node = onnx.helper.make_node( "DequantizeLinear", dq_inputs, [node.output[0]], node.name + "_DequantizeLinear" ) self.quantizer.new_nodes.append(dq_node) # int8 weight will be recalculated for the first time if ( any([child.op_type == "QuantizeLinear" for child in children]) and self.quantizer.model.get_initializer(parents[0].input[0]) is not None and parents[0].input[0] not in self.quantizer.recalculate_quantized_value ): int8_tensor = onnx.numpy_helper.to_array(self.quantizer.model.get_initializer(parents[0].input[0])) in_scale = onnx.numpy_helper.to_array(self.quantizer.model.get_initializer(parents[0].input[1])) in_zp = onnx.numpy_helper.to_array(self.quantizer.model.get_initializer(parents[0].input[2])) new_int8_tensor = (((int8_tensor.astype("float32") - in_zp) * in_scale) / out_scale).round() + out_zp self.quantizer.model.set_initializer(parents[0].input[0], new_int8_tensor.astype(int8_tensor.dtype)) self.quantizer.recalculate_quantized_value.append(parents[0].input[0]) self.quantizer.remove_nodes.extend([node, parents[0]]) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/gavgpool.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """GlobalAveragePool Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="GlobalAveragePool", mode=[constants.STATIC_QUANT]) class GlobalAveragePoolOperator(base_op.Operator): """GlobalAveragePool Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(GlobalAveragePoolOperator, self).__init__(onnx_quantizer, onnx_node) def convert_check(self): """Check if conversion can be done.""" node = self.node children = self.quantizer.model.get_children(node) if len(children) == 0: # pragma: no cover return False return True def convert(self): """Convert to QOperator format.""" node = self.node parent = self.quantizer.model.get_parents(node)[0] child = self.quantizer.model.get_children(node)[0] kwargs = {} for attribute in node.attribute: kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain kwargs["channels_last"] = 0 inputs = parent.input inputs.extend(child.input[1:]) qnode = onnx.helper.make_node("QLinear" + node.op_type, inputs, child.output, node.name + "_quant", **kwargs) self.quantizer.new_nodes += [qnode] self.quantizer.remove_nodes.append(child) self.quantizer.remove_nodes.append(parent) self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/gemm.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Gemm Operator.""" import onnx from onnx_neural_compressor import constants, logger from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Gemm", mode=[constants.STATIC_QUANT]) class GemmOperator(base_op.Operator): """Gemm Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(GemmOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if len(node.input) == 3 and not quant_utils.find_by_name(node.input[2], self.quantizer.model.initializer()): logger.warning( "Bias of Gemm node '{}' is not constant. " "Exclude this node can get better performance.".format(node.name) ) if self.quantizer.quant_format != "qdq": return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0]) if self.per_channel and quant_utils.find_by_name(node.input[1], self.quantizer.model.initializer()): self.quantizer.quantize_weights_per_channel( node, [1], self.weight_dtype, self.weight_sym, 0 if quant_utils.is_B_transposed(node) else 1 ) else: self.quantizer.quantize_inputs(node, [1]) if len(node.input) == 3 and quant_utils.find_by_name(node.input[2], self.quantizer.model.initializer()): self.quantizer.quantize_bias_tensor(node) beta_attribute = [attr for attr in node.attribute if attr.name == "beta"] if len(beta_attribute): beta_attribute[0].f = 1.0 if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) qgemm_inputs = [] for parent in parents[:-1]: qgemm_inputs.extend(parent.input) qgemm_inputs.append(parents[-1].input[0]) kwargs = {} for attribute in node.attribute: if attribute.name != "beta": kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qgemm_output = node.output[0] if not self.disable_qdq_for_node_output: child = self.quantizer.model.get_children(node)[0] self.quantizer.remove_nodes.append(child) qgemm_output = child.output[0] qgemm_inputs.extend(child.input[1:]) qgemm_node = onnx.helper.make_node("QGemm", qgemm_inputs, [qgemm_output], node.name, **kwargs) self.quantizer.new_nodes.append(qgemm_node) self.quantizer.remove_nodes.extend(parents) self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/lstm.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """LSTM Operator.""" import numpy import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="LSTM", mode=[constants.DYNAMIC_QUANT]) class LSTMOperator(base_op.Operator): """LSTM Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(LSTMOperator, self).__init__(onnx_quantizer, onnx_node) def quantize(self): """Do quantizaion.""" return def convert_check(self): """Check if conversion can be done.""" node = self.node if not self.quantizer.is_valid_quantize_weight(node.input[1]) or not self.quantizer.is_valid_quantize_weight( node.input[2] ): # pragma: no cover return False model = self.quantizer.model W = model.get_initializer(node.input[1]) R = model.get_initializer(node.input[2]) if len(W.dims) != 3 or len(R.dims) != 3: # pragma: no cover return False return True def convert(self): """Convert to QOperator format.""" node = self.node model = self.quantizer.model W = model.get_initializer(self.node.input[1]) R = model.get_initializer(self.node.input[2]) [W_num_dir, W_4_hidden_size, W_input_size] = W.dims [R_num_dir, R_4_hidden_size, R_hidden_size] = R.dims if self.per_channel: # pragma: no cover del W.dims[0] del R.dims[0] W.dims[0] = W_num_dir * W_4_hidden_size R.dims[0] = R_num_dir * R_4_hidden_size quant_input_weight_tuple = self.quantizer.quantize_weight_per_channel( node.input[1], self.weight_dtype, self.weight_sym, 0 ) quant_recurrent_weight_tuple = self.quantizer.quantize_weight_per_channel( node.input[2], self.weight_dtype, self.weight_sym, 0 ) W_quant_weight = model.get_initializer(quant_input_weight_tuple[0]) R_quant_weight = model.get_initializer(quant_recurrent_weight_tuple[0]) W_quant_array = onnx.numpy_helper.to_array(W_quant_weight) R_quant_array = onnx.numpy_helper.to_array(R_quant_weight) W_quant_array = numpy.reshape(W_quant_array, (W_num_dir, W_4_hidden_size, W_input_size)) R_quant_array = numpy.reshape(R_quant_array, (R_num_dir, R_4_hidden_size, R_hidden_size)) W_quant_array = numpy.transpose(W_quant_array, (0, 2, 1)) R_quant_array = numpy.transpose(R_quant_array, (0, 2, 1)) W_quant_tranposed = onnx.numpy_helper.from_array(W_quant_array, quant_input_weight_tuple[0]) R_quant_tranposed = onnx.numpy_helper.from_array(R_quant_array, quant_recurrent_weight_tuple[0]) model.remove_initializers([W_quant_weight, R_quant_weight]) model.add_initializer(W_quant_tranposed) model.add_initializer(R_quant_tranposed) W_quant_zp = model.get_initializer(quant_input_weight_tuple[1]) R_quant_zp = model.get_initializer(quant_recurrent_weight_tuple[1]) W_quant_scale = model.get_initializer(quant_input_weight_tuple[2]) R_quant_scale = model.get_initializer(quant_recurrent_weight_tuple[2]) if self.per_channel: # pragma: no cover W_quant_zp.dims[:] = [W_num_dir, W_4_hidden_size] R_quant_zp.dims[:] = [R_num_dir, R_4_hidden_size] W_quant_scale.dims[:] = [W_num_dir, W_4_hidden_size] R_quant_scale.dims[:] = [R_num_dir, R_4_hidden_size] inputs = [] input_len = len(node.input) inputs.extend([node.input[0]]) inputs.extend([quant_input_weight_tuple[0], quant_recurrent_weight_tuple[0]]) inputs.extend([node.input[3] if input_len > 3 else ""]) inputs.extend([node.input[4] if input_len > 4 else ""]) inputs.extend([node.input[5] if input_len > 5 else ""]) inputs.extend([node.input[6] if input_len > 6 else ""]) inputs.extend([node.input[7] if input_len > 7 else ""]) inputs.extend( [ quant_input_weight_tuple[2], quant_input_weight_tuple[1], quant_recurrent_weight_tuple[2], quant_recurrent_weight_tuple[1], ] ) kwargs = {} for attribute in node.attribute: if attribute.name == "layout": continue kwarg = quant_utils.attribute_to_kwarg(attribute) kwargs.update(kwarg) quant_lstm_name = node.name + "_quant" quant_lstm_node = onnx.helper.make_node( "DynamicQuantizeLSTM", inputs, node.output, quant_lstm_name, domain="com.microsoft", **kwargs ) self.quantizer.remove_nodes.append(node) self.quantizer.new_nodes.append(quant_lstm_node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/matmul.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """MatMul Operator.""" import onnx from onnx import onnx_pb as onnx_proto from onnx_neural_compressor import constants from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="MatMul", mode=[constants.DYNAMIC_QUANT]) class MatMulOperator(base_op.Operator): """MatMul Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(MatMulOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if not all([self.quantizer.model.get_initializer(i) is None for i in node.input]): return True elif all([i not in self.quantizer.quantized_value_map for i in node.input]): return False else: return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0]) if self.per_channel and quant_utils.find_by_name(node.input[1], self.quantizer.model.initializer()): self.quantizer.quantize_weights_per_channel(node, [1], self.weight_dtype, self.weight_sym, 1) else: self.quantizer.quantize_inputs(node, [1]) node.name = node.name + "_quant" def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) inputs = [] quantized_name = [] scale = [] zp = [] for parent in parents: if parent.op_type == "DequantizeLinear": quantized_name.append(parent.input[0]) else: quantized_name.append(parent.output[0]) if parent.op_type == "DynamicQuantizeLinear": scale.append(parent.output[1]) zp.append(parent.output[2]) else: scale.append(parent.input[1]) zp.append(parent.input[2]) inputs.extend(quantized_name) inputs.extend(zp) matmul_integer_output = node.output[0] + "_output_quantized" matmul_integer_node = onnx.helper.make_node("MatMulInteger", inputs, [matmul_integer_output], node.name) self.quantizer.new_nodes.append(matmul_integer_node) # Add cast operation to cast matmulInteger output to float. cast_op_output = matmul_integer_output + "_cast_output" cast_node = onnx.helper.make_node( "Cast", [matmul_integer_output], [cast_op_output], matmul_integer_output + "_cast", to=onnx_proto.TensorProto.FLOAT, ) self.quantizer.new_nodes.append(cast_node) # Add mul operation to multiply scales of two inputs. scales_mul_op = node.name + "_scales_mul" scales_mul_node = quant_utils.find_by_name(scales_mul_op, self.quantizer.new_nodes) if scales_mul_node is None: scales_mul_node = onnx.helper.make_node("Mul", [scale[0], scale[1]], [scales_mul_op + ":0"], scales_mul_op) self.quantizer.new_nodes.append(scales_mul_node) scales_mul_op_output = scales_mul_node.output[0] # Add mul operation to multiply mul_scales_op result with output of MatMulInteger # and make the output of this node the same as output of original matmul node. output_scale_mul_op = node.name + "_output_scale_mul" self.quantizer.new_nodes.append( onnx.helper.make_node("Mul", [cast_op_output, scales_mul_op_output], [node.output[0]], output_scale_mul_op) ) if parents[1].op_type == "DequantizeLinear": self.quantizer.remove_nodes.append(parents[1]) self.quantizer.remove_nodes.append(node) @base_op.op_registry(op_types="MatMul", mode=[constants.STATIC_QUANT]) class StaticMatMulOperator(MatMulOperator): """MatMul Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(MatMulOperator, self).__init__(onnx_quantizer, onnx_node) def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0]) if self.per_channel and quant_utils.find_by_name(node.input[1], self.quantizer.model.initializer()): self.quantizer.quantize_weights_per_channel(node, [1], self.weight_dtype, self.weight_sym, 1) else: self.quantizer.quantize_inputs(node, [1]) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) if len(self.quantizer.model.get_children(node)) == 0 or not node.name.endswith("_quant"): # pragma: no cover return qlinear_matmul_inputs = [] if self.disable_qdq_for_node_output: for i in range(len(parents[0].input)): qlinear_matmul_inputs.extend([parent.input[i] for parent in parents]) qlinear_matmul_node = onnx.helper.make_node( "MatMulIntegerToFloat", qlinear_matmul_inputs, node.output, node.name, domain="com.microsoft" ) else: # after inserting QDQ, MatMul -> Q-DQ-MatMul-Q-DQ for parent in parents: qlinear_matmul_inputs.extend(parent.input) child = self.quantizer.model.get_children(node)[0] qlinear_matmul_output = child.output[0] qlinear_matmul_inputs.extend(child.input[1:]) qlinear_matmul_node = onnx.helper.make_node( "QLinearMatMul", qlinear_matmul_inputs, [qlinear_matmul_output], node.name ) self.quantizer.remove_nodes.append(child) self.quantizer.new_nodes.append(qlinear_matmul_node) self.quantizer.remove_nodes.append(node) # make sure parent DequantizeLinear of input 0 is not used by other ops if len(self.quantizer.model.get_children(parents[0])) == 1 and not self.quantizer.model.is_graph_output( parents[0].output[0] ): self.quantizer.remove_nodes.extend(parents) else: self.quantizer.remove_nodes.append(parents[1]) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/maxpool.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """MaxPool Operator.""" from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="MaxPool", mode=[constants.STATIC_QUANT]) class MaxPoolOperator(base_op.Operator): """MaxPool Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(MaxPoolOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node # if opset version is less than 12, just no change if self.quantizer.opset_version < 12: # pragma: no cover return False if not self.quantizer.is_valid_quantize_weight(node.input[0]): # pragma: no cover return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(self.node, direct_int8=True) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(self.node, direct_int8=True) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node children = self.quantizer.model.get_children(node) if len(children) == 0 or not node.name.endswith("_quant"): # pragma: no cover return False return True def convert(self): """Convert to QOperator format.""" node = self.node parent = self.quantizer.model.get_parents(node)[0] children = self.quantizer.model.get_children(node) if parent.op_type != "DequantizeLinear" or all( [i.op_type != "QuantizeLinear" for i in children] ): # pragma: no cover return node.input[0] = parent.input[0] node.output[0] = node.output[0].replace("_QuantizeInput", "_quantized") for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) for n in self.quantizer.model.get_children(child): self.quantizer.model.replace_node_input(n, child.output[0], node.output[0]) self.quantizer.remove_nodes.append(parent) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/pad.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Pad Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Pad", mode=[constants.STATIC_QUANT]) class PadOperator(base_op.Operator): """Pad Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(PadOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" # if opset version is less than 11, just no change if self.quantizer.opset_version < 11: # pragma: no cover return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0]) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(node) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node children = self.quantizer.model.get_children(node) if len(children) == 0 or not node.name.endswith("_quant"): # pragma: no cover return False return True def convert(self): """Convert to QOperator format.""" node = self.node parent = self.quantizer.model.get_parents(node)[0] child = self.quantizer.model.get_children(node)[0] kwargs = {} for attribute in node.attribute: kv = quant_utils.attribute_to_kwarg(attribute) kwargs.update(kv) if "mode" not in kwargs or kwargs["mode"] == b"constant": if len(node.input) > 2: # There is 3rd input 'constant_value' zp_tensor = self.quantizer.model.get_initializer(parent.input[2]) scale_tensor = self.quantizer.model.get_initializer(parent.input[1]) padding_constant_initializer = self.quantizer.model.get_initializer(node.input[2]) if padding_constant_initializer is not None: zp_array = onnx.numpy_helper.to_array(zp_tensor) zp_value = zp_array.item() if zp_array.ndim == 0 else zp_array[0] scale_array = onnx.numpy_helper.to_array(scale_tensor) scale_value = scale_array.item() if scale_array.ndim == 0 else scale_array[0] padding_constant_array = onnx.numpy_helper.to_array(padding_constant_initializer) quantized_padding_constant_array = quant_utils.quantize_nparray( onnx.helper.tensor_dtype_to_np_dtype(self.weight_dtype), padding_constant_array, scale_value, zp_value, ) quantized_padding_constant_name = node.input[2] + "_quantized" quantized_padding_constant_initializer = onnx.numpy_helper.from_array( quantized_padding_constant_array, quantized_padding_constant_name ) # Suppose this padding constant initializer only used by the node self.quantizer.model.remove_initializer(padding_constant_initializer) self.quantizer.model.add_initializer(quantized_padding_constant_initializer) node.input[2] = quantized_padding_constant_name else: self.quantizer.quantize_inputs(node, [2], False) node.input[2] = node.input[2] + "_DequantizeLinear" else: # pad zero_point for original zero node.input.extend([parent.input[2]]) # Create an entry for output quantized value node.input[0] = parent.input[0] node.output[0] = child.output[0] self.quantizer.remove_nodes.extend([parent, child]) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/pooling.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """AveragePool Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="AveragePool", mode=[constants.STATIC_QUANT]) class PoolOperator(base_op.Operator): """AveragePool Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(PoolOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if not self.quantizer.is_valid_quantize_weight(node.input[0]): return False return True def quantize(self): """Do quantizaion.""" node = self.node super().quantize() node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if len(children) == 0 or len(parents) == 0 or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if all([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): qlinear_output_name = node.output[0] + "_quantized" inputs = [] inputs.extend(parents[0].input) inputs.extend([i for i in children if i.op_type == "QuantizeLinear"][0].input[1:]) kwargs = {} for attribute in node.attribute: kwargs.update(quant_utils.attribute_to_kwarg(attribute)) kwargs["domain"] = quant_utils.ms_domain qnode = onnx.helper.make_node("QLinear" + node.op_type, inputs, [qlinear_output_name], node.name, **kwargs) self.quantizer.remove_nodes.extend(parents) for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], qnode.output[0]) self.quantizer.new_nodes.append(qnode) self.quantizer.remove_nodes.append(node) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/reduce.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Reduce Operator.""" from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry( op_types="ReduceMean, ReduceLogSum, ReduceLogSumExp, " "ReduceL1, ReduceL2, ReduceProd, ReduceSum, ReduceSumSquare", mode=[constants.STATIC_QUANT], ) class ReduceOperator(base_op.Operator): """Reduce Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ReduceOperator, self).__init__(onnx_quantizer, onnx_node) @base_op.op_registry(op_types="ReduceMax, ReduceMin", mode=[constants.STATIC_QUANT]) class ReduceMinMaxOperator(base_op.Operator): """ReduceMin and ReduceMax Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ReduceMinMaxOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if not self.quantizer.is_valid_quantize_weight(node.input[0]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(self.node, [0], direct_int8=True) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(self.node, direct_int8=True) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if (len(children) == 0 and len(parents) == 0) or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if any([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): for parent in parents: if parent.op_type == "DequantizeLinear": self.node.input[0] = parent.input[0] self.quantizer.remove_nodes.append(parents[0]) break for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], node.output[0] + "_quantized") node.output[0] = node.output[0] + "_quantized" ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/resize.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Resize Operator.""" from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Resize", mode=[constants.STATIC_QUANT]) class ResizeOperator(base_op.Operator): """Resize Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(ResizeOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node # if version is less than 11, just keep this node if self.quantizer.opset_version < 11: return False if not self.quantizer.is_valid_quantize_weight(node.input[0]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0], direct_int8=True) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(self.node, direct_int8=True) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if (len(children) == 0 and len(parents) == 0) or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if any([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): for parent in parents: if parent.op_type == "DequantizeLinear" and parent.output[0] == node.input[0]: self.node.input[0] = parent.input[0] self.quantizer.remove_nodes.append(parent) break for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], node.output[0] + "_quantized") node.output[0] = node.output[0] + "_quantized" ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/split.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Split Operator.""" import onnx from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Split", mode=[constants.STATIC_QUANT]) class SplitOperator(base_op.Operator): """Split Operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(SplitOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node data_found, _, _, _, _ = self.quantizer._get_quantization_params(node.output[0]) if not data_found: return False if not all([self.quantizer.is_valid_quantize_weight(i) for i in node.input]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(node, [0]) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(self.node, direct_int8=True) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parent = self.quantizer.model.get_parents(node)[0] children = self.quantizer.model.get_children(node) if ( parent.op_type != "DequantizeLinear" or len(children) == 0 or not node.name.endswith("_quant") ): # pragma: no cover return False return True def convert(self): """Convert to QOperator format.""" node = self.node parent = self.quantizer.model.get_parents(node)[0] kwargs = {} for attribute in node.attribute: # pragma: no cover kwargs.update(quant_utils.attribute_to_kwarg(attribute)) quantized_input_names = [] quantized_input_names.append(parent.input[0]) if len(node.input) > 1: # pragma: no cover quantized_input_names.extend(node.input[1:]) outputs = [] input_name_to_nodes = self.quantizer.model.input_name_to_nodes() for output in node.output: if output in input_name_to_nodes: child = input_name_to_nodes[output][0] if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) outputs.append(child.output[0]) else: # pragma: no cover outputs.append(output) else: # pragma: no cover outputs.append(output + "_quantized") quantized_node = onnx.helper.make_node(node.op_type, quantized_input_names, outputs, node.name, **kwargs) self.quantizer.new_nodes.append(quantized_node) self.quantizer.remove_nodes.extend([parent, node]) ================================================ FILE: onnx_neural_compressor/algorithms/post_training_quant/operators/unary_op.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Unary operator.""" from onnx_neural_compressor import constants, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.post_training_quant.operators import base_op @base_op.op_registry(op_types="Exp, Log, Round, Sqrt", mode=[constants.STATIC_QUANT]) class UnaryOperator(base_op.Operator): """Unary operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(UnaryOperator, self).__init__(onnx_quantizer, onnx_node) @base_op.op_registry(op_types="Abs, Shrink, Sign", mode=[constants.STATIC_QUANT]) class UnaryDirect8BitOperator(base_op.Operator): """Unary operator.""" def __init__(self, onnx_quantizer, onnx_node): """Initialization.""" super(UnaryDirect8BitOperator, self).__init__(onnx_quantizer, onnx_node) def quantize_check(self): """Check if quantizaion can be done.""" node = self.node if not self.quantizer.is_valid_quantize_weight(node.input[0]): return False return True def quantize(self): """Do quantizaion.""" node = self.node self.quantizer.quantize_inputs(self.node, [0], direct_int8=True) if not self.disable_qdq_for_node_output: self.quantizer.quantize_outputs(self.node, direct_int8=True) node.name = node.name + "_quant" def convert_check(self): """Check if conversion can be done.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if (len(children) == 0 and len(parents) == 0) or not node.name.endswith("_quant"): return False return True def convert(self): """Convert to QOperator format.""" node = self.node parents = self.quantizer.model.get_parents(node) children = self.quantizer.model.get_children(node) if any([i.op_type == "DequantizeLinear" for i in parents]) and any( [i.op_type == "QuantizeLinear" for i in children] ): for parent in parents: if parent.op_type == "DequantizeLinear": self.node.input[0] = parent.input[0] self.quantizer.remove_nodes.append(parents[0]) break for child in children: if child.op_type == "QuantizeLinear": self.quantizer.remove_nodes.append(child) self.quantizer.model.replace_input_of_all_nodes(child.output[0], node.output[0] + "_quantized") node.output[0] = node.output[0] + "_quantized" ================================================ FILE: onnx_neural_compressor/algorithms/smoother/__init__.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ================================================ FILE: onnx_neural_compressor/algorithms/smoother/calibrator.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Calibration for smooth quant.""" import importlib.util import pathlib import sys import tempfile from typing import List import numpy as np import onnx import onnxruntime from onnx_neural_compressor import data_reader, logger, onnx_model, utility class Calibrator: """Dump information for smooth quant.""" def __init__( self, model: onnx_model.ONNXModel, dataloader: data_reader.CalibrationDataReader, iterations: List[int] = [], providers: List[str] = ["CPUExecutionProvider"], **kwargs, ): """Initialize a Calibrator to dump information. Args: model (onnx_model.ONNXModel): onnx_model.ONNXModel object. dataloader (data_reader.CalibrationDataReader): user implemented object to read in and preprocess calibration dataset. iterations (List[int], optional): tensor of which iteration will be collected. Defaults to []. providers (List[str], optional): execution provider for onnxruntime. Defaults to ["CPUExecutionProvider"]. """ self.model_wrapper = model self.dataloader = dataloader self.augmented_model = None self.iterations = iterations self.providers = providers def _check_is_group_conv(self, node): """Check the op is group wised or not(depthwise conv is excluded,return false). Args: node: The op node Returns: Bool: group wised True, otherwise False, depthwise False """ name_to_indices = {} for index, i in enumerate(self.model_wrapper.initializer()): name_to_indices[i.name] = index if node.op_type == "Conv": group = 1 for attr in node.attribute: if hasattr(attr, "name"): if attr.name == "group": group = attr.i break # currently only normal conv and depthwise conv are supported if group > 1: # group conv, need to check depthwise or not weight_name = node.input[1] weight_shape = onnx.numpy_helper.to_array( self.model_wrapper.initializer()[name_to_indices[weight_name]] ).shape input_channel = weight_shape[1] if input_channel != 1: # TODO: need to double check return True return False def _get_input_tensor_of_ops(self, op_types: List[str] = ["MatMul", "Gemm", "Conv", "FusedConv"]): """Traverse the graph and get all the data tensors flowing into layers of {op_types}. Group conv is excluded. # TODO: the tensors could be set/filtered in configuration. Args: op_types (List[str], optional): The op types whose input tensor will be dumped. Defaults to ["MatMul", "Gemm", "Conv", "FusedConv"]. Returns: dict: A dict of dumped tensor to node info """ tensors_to_node = {} initializers = {i.name: i for i in self.model_wrapper.initializer()} for node in self.model_wrapper.nodes(): if len(op_types) == 0 or node.op_type in op_types: if node.op_type in ["Conv", "FusedConv"] and self._check_is_group_conv(node): continue # also need to check whether the layer has weight if len(node.input) >= 2 and node.input[1] in initializers.keys(): tensors_to_node.setdefault(node.input[0], []).append([node.name, node.input, node.output]) return tensors_to_node def _get_max_per_channel(self, datas, percentile): """Get the max values per input channel. Args: datas: The tensors percentile: percentile of calibration to remove outliers Returns: The max values per input channel """ permute_datas = [] for data in datas: if len(data.shape) == 3: # TODO: mammul batchsize*seq*inchannel, conv:batchsize*inchannle*f*f tensor = np.abs(np.reshape(data, (-1, data.shape[-1]))) permute_datas.append(tensor) elif len(data.shape) == 4: tensor = np.swapaxes(data, 1, -1) tensor = np.abs(np.reshape(tensor, (-1, tensor.shape[-1]))) permute_datas.append(tensor) elif len(data.shape) == 2: permute_datas.append(np.abs(data)) else: assert False, "not supported" permute_datas = np.stack(permute_datas, axis=0) permute_datas = permute_datas.reshape(-1, permute_datas.shape[-1]) max_per_channels = np.percentile(permute_datas, percentile, axis=0) max_per_channels = max_per_channels.astype(np.single) return max_per_channels def get_intermediate_outputs(self): so = onnxruntime.SessionOptions() if sys.version_info < (3, 11) and importlib.util.find_spec("onnxruntime_extensions"): # pragma: no cover from onnxruntime_extensions import get_library_path so.register_custom_ops_library(get_library_path()) providers = self.providers if "TensorrtExecutionProvider" not in self.providers else ["CUDAExecutionProvider"] if self.model_wrapper.is_large_model: # pragma: no cover with tempfile.TemporaryDirectory(prefix="ort.calib.") as tmp_dir: onnx.save_model( self.model_wrapper.model, pathlib.Path(tmp_dir).joinpath("augment.onnx").as_posix(), save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) session = onnxruntime.InferenceSession( pathlib.Path(tmp_dir).joinpath("augment.onnx").as_posix(), so, providers=providers ) onnx.external_data_helper.load_external_data_for_model( self.model_wrapper.model, pathlib.Path(tmp_dir).as_posix() ) else: session = onnxruntime.InferenceSession( self.model_wrapper.model.SerializeToString(), so, providers=providers ) node_output_names = [output.name for output in session.get_outputs()] output_dicts = {} input_name_to_nodes = self.model_wrapper.input_name_to_nodes() output_name_to_node = self.model_wrapper.output_name_to_node() name_to_node = {} for data_name in node_output_names: node = None if data_name in output_name_to_node: node = output_name_to_node[data_name] elif data_name in input_name_to_nodes: node = input_name_to_nodes[data_name][0] assert node, "{} is neither an input nor an output of nodes in augmented model.".format(data_name) name_to_node[data_name] = node.name def _collect_data(ort_inputs): for output_idx, output in enumerate(session.run(None, ort_inputs)): output_dicts.setdefault(node_output_names[output_idx], []).append(output) idx = 0 while True: inputs = self.dataloader.get_next() if not inputs: break if self.iterations != []: if idx > max(self.iterations): break if idx in self.iterations: _collect_data(inputs) else: _collect_data(inputs) idx += 1 return output_dicts def calib_smooth(self, op_types, percentile: float = 99.999): """Smooth model calibration. Mainly get the max info per channel of input tensors. Args: op_types (_type_): The op types whose input tensor will be dumped. percentile (float, optional): Percentile of calibration to remove outliers. Defaults to 99.999. Returns: max_vals_per_channel: max values per channel of input tensors shape_infos: The shape information of input tensors """ logger.info("Start smooth model calibration.") # add the input tensors of {op_types} to outputs of the model tensors_to_node = self._get_input_tensor_of_ops(op_types) self.model_wrapper.add_tensors_to_outputs(tensors_to_node.keys()) output_dicts = self.get_intermediate_outputs() # remove the input tensors of {op_types} to outputs of the model self.model_wrapper.remove_tensors_from_outputs(tensors_to_node.keys()) max_vals_per_channel = {} shape_infos = {} for key, val in tensors_to_node.items(): max_val_per_channel = self._get_max_per_channel(output_dicts[key], percentile=percentile) max_vals_per_channel[key] = max_val_per_channel shape_infos[key] = output_dicts[key][0].shape for item in val: shape_infos[item[1][1]] = self.model_wrapper.get_initializer(item[1][1]).dims return max_vals_per_channel, shape_infos, tensors_to_node ================================================ FILE: onnx_neural_compressor/algorithms/smoother/core.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Smoother for onnxrt.""" import copy import os import pathlib import numpy as np import onnx import onnxruntime as ort from onnx_neural_compressor import data_reader, logger, onnx_model, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.smoother import calibrator from typing import List, Union # isort: skip def _get_quant_dequant_output(model, input_data, output_data, providers): """Get loss between fp32 output and QDQ output. Args: model (object): model input_data (numpy.ndarray): fp32 input output_data (numpy.ndarray): fp32 output providers (list): execution provider """ input_data = quant_utils.qdq_data(input_data, 2, False) sess = ort.InferenceSession(model.SerializeToString(), providers=providers) preds = sess.run(None, {model.graph.input[0].name: input_data}) loss = np.sum(np.abs(output_data - preds) ** 2) return loss def _make_sub_graph(node, inits, input_data, output_data, opset, ir_version): """Build a model with the specific node. Args: node (object): node inits (list): initializer inputs of this node input_data (numpy.ndarray): fp32 input output_data (numpy.ndarray): fp32 output opset (object): opset of the model ir_version (object): ir_version of the model """ input = onnx.helper.make_tensor_value_info( node.input[0], onnx.helper.np_dtype_to_tensor_dtype(input_data.dtype), input_data.shape, ) output = onnx.helper.make_tensor_value_info( node.output[0], onnx.helper.np_dtype_to_tensor_dtype(output_data.dtype), output_data.shape, ) graph = onnx.helper.make_graph([node], "sub_graph", [input], [output], inits) model = onnx.helper.make_model(graph, opset_imports=opset) model.ir_version = ir_version return model class Smoother: """Fake input channel quantization. For more details please refer to: [1] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models [2] SPIQ: Data-Free Per-Channel Static Input Quantization We only support inplace mode which means the model weights will be changed, you can call recover function to recover the weights if needed. """ def __init__( self, model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], dataloader: data_reader.CalibrationDataReader, execution_provider: str = "CPUExecutionProvider", ): """Initialize the attributes of class.""" self.model = ( model if isinstance(model, onnx_model.ONNXModel) else onnx_model.ONNXModel(model, load_external_data=True) ) self.value_infos = {vi.name: vi for vi in self.model.model.graph.value_info} self.value_infos.update({ot.name: ot for ot in self.model.model.graph.output}) self.value_infos.update({it.name: it for it in self.model.model.graph.input}) self.dataloader = dataloader self.providers = [execution_provider] self.tensor_scales_info = {} self.new_added_mul_nodes = [] self.new_added_value_info = [] self.new_init_tensors = [] # scales_tensor self.scales_per_op = True self.replace_input = [] self.ops_to_absorb = [] self.max_vals_per_channel = None self.shape_info = None self.tensors_to_node = None self._build_absorb_function() def transform( self, alpha: Union[float, str] = 0.5, folding: bool = True, percentile: float = 99.999, op_types: List[str] = ["Gemm", "Conv", "MatMul", "FusedConv"], scales_per_op: bool = True, calib_iter: int = 100, auto_alpha_args: dict = {"alpha_min": 0.3, "alpha_max": 0.7, "alpha_step": 0.05, "attn_method": "min"}, *args, **kwargs ): """The main entry of smooth quant. Args: alpha (float, optional): alpha value to balance the quantization difficulty of activation and weight. Defaults to 0.5. folding (bool, optional): whether fold those foldable Mul which are inserted for smooth quant. Defaults to True. percentile (float, optional): percentile of calibration to remove outliers. Defaults to 99.999. op_types (list, optional): the op type to be smooth quantized. Defaults to ["Gemm", "Conv", "MatMul", "FusedConv"]. scales_per_op (bool, optional): True, each op will have an individual scale, mainlyfor accuracy False, ops with the same input will share a scale, mainly for performance. Defaults to True. calib_iter (int, optional): iteration num for calibration. Defaults to 100. auto_alpha_args (_type_, optional): alpha args for auto smooth. Defaults to {"alpha_min": 0.3, "alpha_max": 0.7, "alpha_step": 0.05, "attn_method": "min"}. Returns: onnx.ModelProto: A FP32 model with the same architecture as the orig model but with different weight which will be benefit to quantization """ self.scales_per_op = scales_per_op self.clean() if isinstance(alpha, float) and (alpha < 0 or alpha > 1): logger.warning("alpha should be a float value in [0, 1] or 'auto' ") if alpha < 0: alpha = 0 logger.warning("reset alpha to 0 ") elif alpha > 1.0: alpha = 1.0 logger.warning("reset alpha to 1.0 ") self._dump_op_info(percentile, op_types, calib_iter) if alpha == "auto": alpha = self._auto_tune_alpha(calib_iter, **auto_alpha_args) scales = self._get_smooth_scales(alpha) self._insert_smooth_mul_op(scales) self._adjust_weights(scales) self.model.add_nodes(self.new_added_mul_nodes) self.model.model.graph.value_info.extend(self.new_added_value_info) self.model.add_initializers(self.new_init_tensors) for node, old_input_name, new_input_name in self.replace_input: self.model.replace_node_input(node, old_input_name, new_input_name) self.model.update() if folding: self._fold_scale(scales) self.model.topological_sort() self.model.remove_unused_nodes() return self.model.model def _dump_op_info(self, percentile, op_types, iterations): """Dump op info for smooth quant. Args: percentile (float): percentile of calibration to remove outliers op_types (list): the op type to be smooth quantized iterations (int): iterations """ sq_calibrator = calibrator.Calibrator( self.model, self.dataloader, iterations=list(range(0, iterations)), execution_provider=self.providers, ) self.max_vals_per_channel, self.shape_info, self.tensors_to_node = sq_calibrator.calib_smooth( op_types, percentile ) for node in self.model.nodes(): for out in node.output: if ( out in self.tensors_to_node and node.op_type in self.could_absorb_optype and self.model.get_initializer(node.input[1]) is not None ): self.ops_to_absorb.append(node.name) def recover(self): """Recover the model weights.""" for tensor_name, nodes in self.tensors_to_node.items(): for node_info in nodes: key = node_info[0] if self.scales_per_op else tensor_name if key not in self.tensor_scales_info: continue input = node_info[1][1] weight = onnx.numpy_helper.to_array( self.model.get_initializer(input), base_dir=os.path.dirname(self.model.model_path) if self.model.model_path is not None else "", ) scale = self.tensor_scales_info[key] new_weight = weight * scale self.model.set_initializer(input, new_weight) for node, old_input_name, new_input_name in self.replace_input: self.model.replace_node_input(node, new_input_name, old_input_name) for value_info in self.new_added_value_info: self.model.model.graph.value_info.remove(value_info) self.model.remove_nodes(self.new_added_mul_nodes) self.model.remove_initializers(self.new_init_tensors) self.tensor_scales_info = {} self.new_added_mul_nodes = [] self.new_init_tensors = [] self.new_added_value_info = [] self.replace_input = [] def clean(self): """Clean data collected from calibration.""" self.tensor_scales_info = {} self.new_added_mul_nodes = [] self.new_init_tensors = [] self.new_added_value_info = [] self.replace_input = [] def _build_absorb_function(self): """Build function mapping for scale folding.""" def norm(node, scale): # pragma: no cover for idx in [1, 2]: tensor = self.model.get_initializer(node.input[idx]) new_tensor = ( onnx.numpy_helper.to_array(tensor, os.path.dirname(self.model.model_path)) * scale if self.model.model_path is not None else onnx.numpy_helper.to_array(tensor) * scale ) self.model.set_initializer(node.input[idx], new_tensor) self.tensor_scales_info[node.input[idx]] = ( 1.0 / scale if node.input[idx] not in self.tensor_scales_info else self.tensor_scales_info[node.input[idx]] * 1.0 / scale ) return True def mul(node, scale): # pragma: no cover if all([self.model.get_initializer(inp) is None for inp in node.input]): return False for inp in node.input: if self.model.get_initializer(inp) is not None: key = node.input[0].split("_smooth_output")[0] tensor = self.model.get_initializer(inp) new_tensor = ( onnx.numpy_helper.to_array(tensor, os.path.dirname(self.model.model_path)) * scale if self.model.model_path is not None else onnx.numpy_helper.to_array(tensor) * scale ) # set_initializer requires the dims of old & new initializers are same # Mul operator has broadcast mechanism self.model.remove_initializer(tensor) self.model.add_initializer( onnx.helper.make_tensor( inp, tensor.data_type, list(new_tensor.shape), new_tensor.flatten().tolist() ) ) self.tensor_scales_info[key] = ( 1.0 / scale if key not in self.tensor_scales_info else 1.0 / scale * self.tensor_scales_info[key] ) return True def conv(node, scale): # pragma: no cover if len(node.input) > 2: if self.model.get_initializer(node.input[2]) is not None: tensor = self.model.get_initializer(node.input[2]) new_tensor = ( onnx.numpy_helper.to_array(tensor, os.path.dirname(self.model.model_path)) * scale if self.model.model_path is not None else onnx.numpy_helper.to_array(tensor) * scale ) self.model.set_initializer(node.input[2], new_tensor) self.tensor_scales_info[node.input[2]] = 1.0 / scale scale = scale.reshape(-1, 1, 1, 1) tensor = self.model.get_initializer(node.input[1]) new_tensor = ( onnx.numpy_helper.to_array(tensor, os.path.dirname(self.model.model_path)) * scale if self.model.model_path is not None else onnx.numpy_helper.to_array(tensor) * scale ) self.model.set_initializer(node.input[1], new_tensor) self.tensor_scales_info[node.input[1]] = ( 1.0 / scale if node.input[1] not in self.tensor_scales_info else self.tensor_scales_info[node.input[1]] * 1.0 / scale ) return True self.could_absorb_optype = { "LayerNormalization": norm, "BatchNormalization": norm, "InstanceNormalization": norm, "SimplifiedLayerNormalization": mul, "MatMul": mul, "Gemm": mul, "Conv": conv, "FusedConv": conv, "Mul": mul, } def _fold_scale(self, scales): """Absorb the scale to the operator at output channel. Args: scales (dict): scales for smooth quant, {tensor_name: smooth quant scale} """ remove_nodes = [] for node in self.model.nodes(): if node.op_type == "Mul" and node.name.endswith("_smooth_mul") and node not in remove_nodes: parent = self.model.get_parent(node, 0) if parent is None: continue if parent.op_type in self.could_absorb_optype and len(self.model.get_children(parent)) == 1: if node.output[0].split("_smooth_output")[0] in scales: if self.could_absorb_optype[parent.op_type]( parent, 1.0 / scales[node.output[0].split("_smooth_output")[0]] ): remove_nodes.append(node) children = [i for i in self.model.nodes() if node.output[0] in i.input] for child in children: for idx, inp in enumerate(child.input): if inp == node.output[0]: child.input[idx] = node.input[0] self.model.remove_nodes(remove_nodes) def _get_output_loss(self, node_name, scale, calib_iter): """Get output loss of specific node after inserting QDQ pair. Args: node_name (str): node name scale (float): scale of the specific node calib_iter (int): iterations """ node = [i for i in self.model.nodes() if i.name == node_name] loss = 0 if len(node) > 0: node = node[0] orig_outputs = self.model.output() added_tensors = [node.input[0], node.output[0]] self.model.add_tensors_to_outputs(added_tensors) session = ( ort.InferenceSession(self.model.model_path + "_augment.onnx", providers=self.providers) if self.model.is_large_model else ort.InferenceSession(self.model.model.SerializeToString(), providers=self.providers) ) base_dir = "" if not self.model.is_large_model else os.path.dirname(self.model.model_path) weight = onnx.numpy_helper.to_array(self.model.get_initializer(node.input[1]), base_dir) weight_q = quant_utils.qdq_data(weight, 3, True) self.model.set_initializer(node.input[1], weight_q) inits = [self.model.get_initializer(i) for i in node.input if self.model.get_initializer(i) is not None] model = None idx = 1 while True: inputs = self.dataloader.get_next() if not inputs: break if idx > calib_iter: break outputs = session.run(added_tensors, inputs) if model is None: model = _make_sub_graph( node, inits, outputs[0], outputs[1], self.model.model.opset_import, self.model.model.ir_version, ) loss += _get_quant_dequant_output(model, outputs[0] * scale, outputs[1], self.providers) self.model.remove_tensors_from_outputs([i for i in added_tensors if i not in orig_outputs]) self.model.set_initializer(node.input[1], weight) return loss def _reshape_scale_for_input(self, tensor, key): """Reshape the scale for input feature in channel. Args: tensor (str): tensor name key (str): scale key of this tensor """ if len(self.shape_info[tensor]) == 4: scale = np.reshape(self.tensor_scales_info[key], (1, self.tensor_scales_info[key].shape[1], 1, 1)) else: scale = np.reshape(self.tensor_scales_info[key], (1, self.tensor_scales_info[key].shape[0])) return scale def _auto_tune_alpha( self, calib_iter, alpha_min: float = 0.3, alpha_max: float = 0.7, alpha_step: float = 0.05, attn_method: str = "min", ): """Perform alpha-tuning to obtain layer-wise optimal alpha values and adjust parameters accordingly. Args: calib_iter (int): iterations alpha_min (float): min value of alpha search space. alpha_max (float): max value of alpha search space. alpha_step (float): step size of alpha search space. attn_method (str): criterion method used on attention ops; currently min, max and mean are supported. """ logger.info("auto tuning alpha") alpha_space = np.arange(alpha_min, alpha_max, alpha_step).tolist() optimal_alphas = {} if self.model.is_large_model: onnx.save_model( self.model.model, self.model.model_path + "_augment.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="weights.pb", convert_attribute=False, ) ## Searching optimal alphas for tensor_name, node_infos in self.tensors_to_node.items(): for node_info in node_infos: loss_alpha = {} key = node_info[0] if self.scales_per_op else tensor_name node = self.model.get_node(node_info[0]) for alpha in alpha_space: scale = self._get_smooth_scales(alpha, [key]) self._adjust_weights(scale) input_scale = ( self._reshape_scale_for_input(tensor_name, key) if not (node.op_type == "Gemm" and quant_utils.is_B_transposed(node)) else self.tensor_scales_info[key] ) loss = self._get_output_loss(node_info[0], input_scale, calib_iter) loss_alpha[alpha] = loss if key not in optimal_alphas: # Update alpha results optimal_alphas[key] = alpha else: optimal_alphas[key] = ( alpha if optimal_alphas[key] in loss_alpha and loss < loss_alpha[optimal_alphas[key]] else optimal_alphas[key] ) self.recover() logger.info("auto tuning alpha done") if self.model.is_large_model: onnx.external_data_helper.load_external_data_for_model( self.model.model, os.path.split(self.model.model_path)[0] ) os.remove(self.model.model_path + "_augment.onnx") os.remove(os.path.join(os.path.dirname(self.model.model_path), "weights.pb")) return optimal_alphas def _get_smooth_scales(self, alpha, target_list=[]): """Get the smooth scales for. The ops with the same input will share one mul layer. TODO support individual scales for each layer. Args: alpha: smooth alpha in paper target_list: target objects to get scale, [] means get all scales Returns: the smooth scales for weights, currently one input tensor only have one scale """ scales = {} for tensor, nodes in self.tensors_to_node.items(): # if scales_per_op the key of scales is the node name, otherwise the activation of node if self.scales_per_op: for node_info in nodes: node = self.model.get_node_by_weight(node_info[1][1]) if len(target_list) > 0 and node_info[0] not in target_list: continue weight = onnx.numpy_helper.to_array( self.model.get_initializer(node_info[1][1]), base_dir=os.path.dirname(self.model.model_path) if self.model.model_path is not None else "", ) if (len(weight.shape) == 4 and weight.shape[1] != 1) or ( node.op_type == "Gemm" and quant_utils.is_B_transposed(node) ): weight = np.moveaxis(weight, 0, 1) specific_alpha = alpha[node_info[0]] if isinstance(alpha, dict) else alpha scales[node_info[0]] = self._get_smooth_scale(weight, specific_alpha, tensor) else: if len(target_list) > 0 and tensor not in target_list: continue weights_in_channel_max = [] for node_info in nodes: node = self.model.get_node_by_weight(node_info[1][1]) weight = onnx.numpy_helper.to_array( self.model.get_initializer(node_info[1][1]), base_dir=os.path.dirname(self.model.model_path) if self.model.model_path is not None else "", ) if (len(weight.shape) == 4 and weight.shape[1] != 1) or ( node.op_type == "Gemm" and quant_utils.is_B_transposed(node) ): weight = np.moveaxis(weight, 0, 1) weight = weight.reshape(weight.shape[0], -1) cur_max = np.amax(weight, axis=-1) weights_in_channel_max.append(cur_max) weights_stack = np.stack(weights_in_channel_max, axis=-1) specific_alpha = alpha[tensor] if isinstance(alpha, dict) else alpha scales[tensor] = self._get_smooth_scale(weights_stack, specific_alpha, tensor) return scales def _get_smooth_scale(self, weights, specific_alpha, tensor): """Get smooth scale for specific weight. Args: weights (numpy.ndarray): weight data specific_alpha (float): current alpha for this weights tensor (str): tensor name """ weights = np.abs(weights.reshape(weights.shape[0], -1)) weights_max = np.amax(weights, axis=-1) input_power = np.power(self.max_vals_per_channel[tensor], specific_alpha) weight_power = np.power(weights_max, 1 - specific_alpha) weight_power = np.clip(weight_power, a_min=1e-5, a_max=None) scale = np.clip(input_power / weight_power, a_min=1e-5, a_max=None) return scale def _insert_smooth_mul_op(self, scales): """Insert the Mul after inupt. The ops with the same input will share one mul layer. Args: scales (dict): The smooth scales """ for key in scales.keys(): input_name = key if not self.scales_per_op else self.model.get_node(key).input[0] weight_name = ( self.tensors_to_node[key][0][1][1] if not self.scales_per_op else self.model.get_node(key).input[1] ) scale_factor = 1.0 / scales[key] if ( len(self.shape_info[weight_name]) == 3 or len(self.shape_info[weight_name]) == 2 ): # the last dim is input channel pass elif len(self.shape_info[weight_name]) == 4: scale_factor = np.reshape(scale_factor, (1, -1, 1, 1)) else: assert False, "not support" name = key + "_" + "smooth_scale" scale_tensor = onnx.helper.make_tensor( name=key + "_" + "smooth_scale", data_type=onnx.TensorProto.FLOAT, dims=scale_factor.shape, vals=scale_factor.flatten().tolist(), ) self.new_init_tensors.append(scale_tensor) mul_output_name = key + "_smooth_output" mul_node = onnx.helper.make_node( "Mul", inputs=[input_name, key + "_" + "smooth_scale"], outputs=[mul_output_name], name=key + "_smooth_mul", ) self.new_added_mul_nodes.append(mul_node) if input_name in self.value_infos: value_info = copy.deepcopy(self.value_infos[input_name]) value_info.name = mul_node.output[0] self.new_added_value_info.append(value_info) if self.scales_per_op: self.replace_input.append([self.model.get_node(key), input_name, mul_output_name]) else: for node_info in self.tensors_to_node[key]: self.replace_input.append([self.model.get_node(node_info[0]), key, mul_output_name]) def _adjust_weights(self, scales): """Adjust the weights with scale. Args: scales (dict): The input scales """ for idx, (tensor_name, nodes) in enumerate(self.tensors_to_node.items()): utility.simple_progress_bar(len(self.tensors_to_node), idx + 1) for node_info in nodes: key = node_info[0] if self.scales_per_op else tensor_name if key not in scales: continue input = node_info[1][1] node = self.model.get_node_by_weight(input) weight = onnx.numpy_helper.to_array( self.model.get_initializer(input), base_dir=os.path.dirname(self.model.model_path) if self.model.model_path is not None else "", ) if len(weight.shape) == 2: scale = ( np.expand_dims(scales[key], axis=0) if node.op_type == "Gemm" and quant_utils.is_B_transposed(node) else np.expand_dims(scales[key], axis=-1) ) new_weight = weight * scale elif len(weight.shape) == 4: # TODO need to check conv node = self.model.get_node_by_weight(input) if ( weight.shape[1] == 1 and "group" in [i.name for i in node.attribute] and [i for i in node.attribute if i.name == "group"][0].i > 1 ): scale = np.reshape(scales[key], (-1, 1, 1, 1)) else: scale = np.reshape(scales[key], (1, -1, 1, 1)) new_weight = weight * scale else: assert False, "not support" self.tensor_scales_info[key] = 1.0 / scale new_tensor = onnx.numpy_helper.from_array(new_weight, input) self.model.get_initializer(input).CopyFrom(new_tensor) ================================================ FILE: onnx_neural_compressor/algorithms/weight_only/__init__.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ================================================ FILE: onnx_neural_compressor/algorithms/weight_only/awq.py ================================================ # Copyright (c) 2023 MIT HAN Lab # This source code is licensed under the MIT license # # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import os import pathlib import numpy as np import onnx import onnxruntime as ort from packaging import version from onnx_neural_compressor import constants, data_reader, logger, onnx_model from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.weight_only import rtn from typing import List, Union # isort: skip def _get_weight_scale(weight, group_size): """Get the scale of weight.""" org_shape = weight.shape weight = np.reshape(weight, (-1, group_size)) if group_size != -1 else weight scale = np.mean(np.reshape(np.abs(weight) / np.max(np.abs(weight), axis=1, keepdims=True), org_shape), axis=0) return scale def _apply_awq_scale(model, weight_config, absorb_pairs, output_dicts): """Apply scale for salient weight.""" best_scales = {} new_init_tensors = [] new_added_mul_nodes = [] replace_input = [] updated_nodes = [] base_dir = os.path.dirname(model.model_path) if model.model_path is not None else "" input_name_to_nodes = model.input_name_to_nodes() for parent, nodes in absorb_pairs.items(): if any([node.input[0] not in output_dicts for node in nodes]): # pragma: no cover logger.warning( "Miss input tensors of nodes {} during AWQ, skip it!".format( ", ".join([node.name for node in nodes if node.input[0] not in output_dicts]) ) ) continue inp = np.concatenate(output_dicts[nodes[0].input[0]], axis=0) inp_scale = np.mean(np.reshape(np.abs(inp), (-1, inp[0].shape[-1])), axis=0) dtype = None weight = [] org_out = [] weight_dtype = weight_config[nodes[0].name].get("weight_dtype", "int") num_bits = weight_config[nodes[0].name].get("weight_bits", 4) group_size = weight_config[nodes[0].name].get("weight_group_size", 32) sym = weight_config[nodes[0].name].get("weight_sym", True) accuracy_level = weight_config[nodes[0].name].get("accuracy_level", 0) # use same params for all children of one parent for node in nodes: weight_config.setdefault(node.name, {}).update({"weight_dtype": weight_dtype}) weight_config.setdefault(node.name, {}).update({"weight_bits": num_bits}) weight_config.setdefault(node.name, {}).update({"weight_group_size": group_size}) weight_config.setdefault(node.name, {}).update({"weight_sym": sym}) # search scale best_error = float("inf") best_ratio = -1 best_scale = None n_grid = 20 for ratio in range(n_grid): ratio = ratio * 1 / n_grid loss = 0 for node in nodes: weight = onnx.numpy_helper.to_array(model.get_initializer(node.input[1]), base_dir) if len(weight.shape) != 2: continue org_out = np.matmul(inp, weight) org_w_shape = weight.shape group_size = group_size if group_size != -1 else org_w_shape[0] w_scale = _get_weight_scale(weight.T, weight.shape[0]) scales = np.clip(np.power(inp_scale, ratio) / np.power(w_scale, (1 - ratio)), 1e-4, None) scales = scales / np.sqrt(np.max(scales) * np.min(scales)) weight = weight.T * scales weight = quant_utils.pad_tensor(weight.T, group_size, (org_w_shape[0] + group_size - 1) // group_size) q_weight = quant_utils.qdq_data( weight.reshape((-1, group_size)), weight_dtype + str(num_bits), sym, ).reshape(weight.shape) q_weight = q_weight[: org_w_shape[0], :] / np.expand_dims(scales, axis=-1) out = np.matmul(inp, q_weight) loss += np.mean(np.power((org_out - out), 2)) is_best = loss < best_error if is_best: best_error = loss best_ratio = ratio best_scale = scales for node in nodes: init_share_num = model.get_initializer_share_num(node.input[1]) weight_tensor = model.get_initializer(node.input[1]) tensor = onnx.numpy_helper.to_array(weight_tensor, base_dir) dtype = tensor.dtype tensor = tensor.T * best_scale tensor = (tensor.T).astype(dtype) new_tensor = onnx.helper.make_tensor( name=node.input[1] + "_scaled", data_type=onnx.helper.np_dtype_to_tensor_dtype(dtype), dims=tensor.shape, vals=tensor.tobytes(), raw=True, ) model.add_initializer(new_tensor) node.input[1] = new_tensor.name if init_share_num == 1: model.remove_initializer(weight_tensor) if parent is None: continue parent = model.get_node(parent) if parent is None or parent.name in updated_nodes: continue if parent.op_type in ["LayerNormalization", "BatchNormalization", "InstanceNormalization"] and len( input_name_to_nodes[nodes[0].input[0]] ) == len( nodes ): # pragma: no cover for idx in [1, 2]: tensor = onnx.numpy_helper.to_array(model.get_initializer(parent.input[idx]), base_dir) dtype = tensor.dtype new_tensor = tensor / np.reshape(best_scale, (1, -1)) model.set_initializer(parent.input[idx], new_tensor.astype(dtype), raw=True) updated_nodes.append(parent.name) output_dicts[parent.output[0]] = output_dicts[parent.output[0]] / np.reshape(best_scale, (1, -1)) elif ( parent.op_type in ["SimplifiedLayerNormalization", "MatMul", "Gemm", "Mul"] and not all([model.get_initializer(inp) is None for inp in parent.input]) and len(input_name_to_nodes[nodes[0].input[0]]) == len(nodes) ): # pragma: no cover for inp in parent.input: if model.get_initializer(inp) is not None: tensor = onnx.numpy_helper.to_array(model.get_initializer(inp), base_dir) dtype = tensor.dtype new_tensor = tensor / np.reshape(best_scale, (1, -1)) model.set_initializer(inp, new_tensor.astype(dtype), raw=True) updated_nodes.append(parent.name) output_dicts[parent.output[0]] = output_dicts[parent.output[0]] / np.reshape(best_scale, (1, -1)) elif parent.op_type in ["Conv", "FusedConv"] and len(input_name_to_nodes[nodes[0].input[0]]) == len( nodes ): # pragma: no cover tensor = onnx.numpy_helper.to_array(model.get_initializer(parent.input[2]), base_dir) dtype = tensor.dtype new_tensor = tensor / np.reshape(best_scale, (1, -1)) model.set_initializer(parent.input[2], new_tensor.astype(dtype), raw=True) updated_nodes.append(parent.name) output_dicts[parent.output[0]] = output_dicts[parent.output[0]] / np.reshape(best_scale, (1, -1)) else: # insert mul scale_tensor = onnx.helper.make_tensor( name=parent.output[0] + "_weight_only_scale", data_type=onnx.helper.np_dtype_to_tensor_dtype(dtype), dims=best_scale.shape, vals=(1.0 / best_scale).flatten().tolist(), ) new_init_tensors.append(scale_tensor) mul_output_name = parent.output[0] + "_weight_only_out" mul_node = onnx.helper.make_node( "Mul", inputs=[nodes[0].input[0], scale_tensor.name], outputs=[mul_output_name], name=nodes[0].input[0] + "_weight_only_mul", ) new_added_mul_nodes.append(mul_node) for node in nodes: replace_input.append([node, node.input[0], mul_node.output[0]]) updated_nodes.append(parent.name) output_dicts[mul_node.output[0]] = output_dicts[mul_node.input[0]] / np.reshape(best_scale, (1, -1)) model.add_nodes(new_added_mul_nodes) model.add_initializers(new_init_tensors) for node, old_input_name, new_input_name in replace_input: model.replace_node_input(node, old_input_name, new_input_name) return model, output_dicts def _apply_awq_clip(model, weight_config, absorb_pairs, output_dicts): """Apply clip for weight by checking mse.""" base_dir = os.path.dirname(model.model_path) if model.model_path is not None else "" ratios = {} for parent, nodes in absorb_pairs.items(): if any([node.input[0] not in output_dicts for node in nodes]): logger.warning( "Miss input tensors of nodes {} during AWQ, skip it!".format( ", ".join([node.name for node in nodes if node.input[0] not in output_dicts]) ) ) continue inp = np.concatenate(output_dicts[nodes[0].input[0]], axis=0) for node in nodes: weight_dtype = weight_config[node.name].get("weight_dtype", "int") num_bits = weight_config[node.name].get("weight_bits", 4) group_size = weight_config[node.name].get("weight_group_size", 32) sym = weight_config[node.name].get("weight_sym", True) accuracy_level = weight_config[node.name].get("accuracy_level", 0) org_weight = onnx.numpy_helper.to_array(model.get_initializer(node.input[1]), base_dir=base_dir) org_w_shape = org_weight.shape # ic, oc group_size = group_size if group_size != -1 else org_w_shape[0] org_out = np.matmul(inp, org_weight) # n_token, oc k_blocks = (org_w_shape[0] - 1) // group_size + 1 org_weight = quant_utils.pad_tensor(org_weight, group_size, k_blocks) org_weight = np.transpose(org_weight) best_error = float("inf") best_ratio = 1 for i_s in range(10): ratio = 1 - i_s / 100 weight = copy.deepcopy(org_weight) weight = quant_utils.qdq_data( weight.reshape((-1, group_size)), weight_dtype + str(num_bits), sym, ratio=ratio, ).reshape(org_weight.shape) cur_out = np.matmul(inp, weight[:, : org_w_shape[0]].T) loss = np.mean(np.power((org_out - cur_out), 2)) is_best = loss < best_error if is_best: best_error = loss best_ratio = ratio ratios[node.input[1]] = best_ratio return ratios def awq_quantize( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], data_reader: data_reader.CalibrationDataReader, weight_config: dict = {}, enable_auto_scale: bool = True, enable_mse_search: bool = True, providers: List[str] = ["CPUExecutionProvider"], ) -> onnx.ModelProto: """Quant the model with Activation-aware Weight quantization(AWQ) method. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): onnx model. data_reader (data_reader.CalibrationDataReader): data_reader for calibration. weight_config (dict, optional): quantization config For example, weight_config = { '(fc2, "MatMul")': { 'weight_dtype': 'int', 'weight_bits': 4, 'weight_group_size': 32, 'weight_sym': True, 'accuracy_level': 0 } }. Defaults to {}. enable_auto_scale (bool, optional): whether to search for best scales based on activation distribution. Defaults to True. enable_mse_search (bool, optional): whether to search for the best clip range from range [0.91, 1.0, 0.01]. Defaults to True. providers (list, optional): providers to use. Defaults to ["CPUExecutionProvider"]. Returns: onnx.ModelProto: quantized onnx model. """ if not isinstance(model, onnx_model.ONNXModel): model = onnx_model.ONNXModel(model) output_dicts = {} full_ratio = {} if enable_mse_search: inputs, so = quant_utils.prepare_inputs(model, data_reader, providers) del data_reader org_output = copy.deepcopy(model.model.graph.output) model.remove_tensors_from_outputs([i.name for i in org_output]) output_names = [] for node in model.nodes(): # check op_type of node is MatMul # check op_name in quantization config # check dim 1 of input is weight tensor if ( node.op_type in ["MatMul"] and node.name in weight_config and model.get_initializer(node.input[1]) is not None ): output_names.append(node.input[0]) output_names = list(set(output_names)) model.add_tensors_to_outputs(output_names) if model.is_large_model: # pragma: no cover onnx.save_model( model.model, model.model_path + "_augment.onnx", save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) session = ( ort.InferenceSession(model.model.SerializeToString(), so, providers=providers) if not model.is_large_model else ort.InferenceSession(model.model_path + "_augment.onnx", so, providers=providers) ) output_name_to_node = model.output_name_to_node() input_name_to_nodes = model.input_name_to_nodes() for input_name in output_names: # input_name maybe the input of graph and there is no parent node parent = output_name_to_node[input_name].name if input_name in output_name_to_node else None dump_pairs = {parent: []} for node in input_name_to_nodes[input_name]: # check op_type of node is MatMul # check op_name in quantization config # check dim 1 of input is weight tensor if ( node.op_type in ["MatMul"] and node.name in weight_config and model.get_initializer(node.input[1]) is not None ): dump_pairs[parent].append(model.get_node(node.name)) if len(dump_pairs[parent]) == 0: # pragma: no cover continue output_dicts = {} for inp in inputs: output = session.run([input_name], inp) output_dicts.setdefault(input_name, []).append(output) if enable_auto_scale: model, output_dicts = _apply_awq_scale( model, weight_config, dump_pairs, output_dicts, ) if enable_mse_search: ratios = _apply_awq_clip( model, weight_config, dump_pairs, output_dicts, ) del output_dicts del dump_pairs full_ratio.update(ratios) model.remove_tensors_from_outputs(output_names) model.model.graph.output.MergeFrom(org_output) model = rtn.rtn_quantize( model=model, weight_config=weight_config, ratios=full_ratio, providers=providers, ) return model def apply_awq_on_model( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], quant_config: dict, calibration_data_reader: data_reader.CalibrationDataReader, enable_auto_scale: bool = True, enable_mse_search: bool = True, providers: List[str] = ["CPUExecutionProvider"], ) -> onnx.ModelProto: """Apply Activation-aware Weight quantization(AWQ) on onnx model. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): nnx model. quant_config (dict): quantization config. calibration_data_reader (data_reader.CalibrationDataReader): data_reader for calibration. Returns: onnx.ModelProto: quantized onnx model. """ # set model params kwargs = { "enable_auto_scale": enable_auto_scale, "enable_mse_search": enable_mse_search, "providers": providers, } q_model = awq_quantize(model, data_reader=calibration_data_reader, weight_config=quant_config, **kwargs) quant_utils.dump_woq_stats(q_model, quant_config) return q_model ================================================ FILE: onnx_neural_compressor/algorithms/weight_only/gptq.py ================================================ # Copyright (c) 2023 MIT HAN Lab # This source code is licensed under the MIT license # # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import os import pathlib import numpy as np import onnx import onnxruntime as ort from onnx_neural_compressor import constants, data_reader, onnx_model, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.layer_wise import core from onnx_neural_compressor.algorithms.weight_only import rtn from onnx_neural_compressor.quantization import config from typing import List, Union # isort: skip def _gptq( W: np.array, H: np.array, num_bits: int = 4, group_size: int = 32, sym: bool = False, block_size: int = 128, percdamp: float = 0.01, actorder: bool = False, mse: bool = False, perchannel: bool = True, ): """Quant the weight with GPTQ method. Args: W (np.array): weight. H (np.array): Hessian matrix. num_bits (int, optional): num_bits. Default is 4. group_size (int, optional): how many elements share one scale/zp. Default is 32. sym (bool, optional): sym or asym. Defaults to False. block_size (int, optional): block_size to quantize weight. percdamp (float, optional): percent of the average Hessian diagonal to use for dampening. actorder (bool, optional): whether rearrange Hessian matrix considering the diag's value. mse (bool, optional): whether get scale and zero point with mse error. perchannel (bool, optional): whether quantize weight per-channel. Returns: Q: fake quantized weight """ Qs = [] maxq = 2**num_bits - 1 grid = 100 maxshrink = 0.8 norm = 2.4 def find_params(weight): org_shape = weight.shape # find zp, scale if not perchannel: weight = np.expand_dims(weight.flatten(), axis=1) tmp = np.zeros(weight.shape[1]) xmin = np.minimum(np.min(weight, axis=0), tmp) xmax = np.maximum(np.max(weight, axis=0), tmp) if sym: xmax = np.maximum(np.abs(xmin), xmax) tmp = xmin < 0 if np.any(tmp): xmin[tmp] = -xmax[tmp] tmp = (xmin == 0) & (xmax == 0) xmin[tmp] = -1 xmax[tmp] = +1 scale = (xmax - xmin) / maxq if sym: zero = np.ones(scale.shape) * (maxq + 1) / 2 else: zero = np.round(-xmin / scale) if mse: best = np.ones([weight.shape[1]]) * float("inf") for i in range(int(maxshrink * grid)): p = 1 - i / grid xmin1 = p * xmin xmax1 = p * xmax scale1 = (xmax1 - xmin1) / maxq zero1 = np.round(-xmin1 / scale1) if not sym else zero q = np.clip(np.round(weight / scale1) + zero1, 0, maxq) q -= weight q = np.power(np.abs(q), norm) err = np.sum(q, 0) tmp = err < best if np.any(tmp): best[tmp] = err[tmp] scale[tmp] = scale1[tmp] zero[tmp] = zero1[tmp] if not perchannel: tmp = org_shape[1] scale = np.repeat(scale, tmp) zero = np.repeat(zero, tmp) shape = [-1] + [1] * (len(org_shape) - 1) scale = np.reshape(scale, shape) zero = np.reshape(zero, shape) return scale, zero scales = [] zps = [] shape = W.shape scale, zp = find_params(W) dead = np.diag(H) == 0 H[dead, dead] = 1 W[dead, :] = 0 # such channel makes no contribution to quantization computation # rearrange considering the diag's value if actorder: perm = np.argsort(np.diag(H))[::-1] W = W[perm, :] H = H[perm, :][:, perm] Losses = np.zeros_like(W) Q = np.zeros_like(W) damp = percdamp * np.mean(np.diag(H)) diag = np.arange(shape[0]) H[diag, diag] += damp # add a average value of H = np.linalg.cholesky(np.linalg.inv(H)).T Hinv = H for i1 in range(0, shape[0], block_size): i2 = min(i1 + block_size, shape[0]) count = i2 - i1 W1 = copy.deepcopy(W[i1:i2, :]) Q1 = np.zeros_like(W1) Err1 = np.zeros_like(W1) Losses1 = np.zeros_like(W1) Hinv1 = Hinv[i1:i2, i1:i2] for i in range(count): # within a block, channel wise w = W1[i, :] d = Hinv1[i, i] if group_size != -1: if (i1 + i) % group_size == 0: scale, zp = find_params(W[(i1 + i) : (i1 + i + group_size), :]) q = (scale * (np.clip(np.round(w[:, np.newaxis] / scale) + zp, 0, maxq) - zp)).flatten() Q1[i, :] = q Losses1[i, :] = (w - q) ** 2 / d**2 err1 = (w - q) / d W1[i:, :] -= np.matmul(np.expand_dims(Hinv1[i:, i], axis=1), np.expand_dims(err1, axis=0)) Err1[i, :] = err1 Q[i1:i2, :] = Q1 Losses[i1:i2, :] = Losses1 / 2 W[i2:, :] -= np.matmul(Hinv[i2:, i1:i2], Err1) if actorder: invperm = np.argsort(perm) Q = Q[invperm, :] Q = np.reshape(Q, W.shape) del W return Q def gptq_quantize( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], data_reader: data_reader.CalibrationDataReader, weight_config: dict = {}, percdamp: float = 0.01, block_size: int = 128, actorder: bool = False, mse: bool = False, perchannel: bool = True, providers: List[str] = ["CPUExecutionProvider"], return_modelproto: bool = True, ): """Quant the model with GPTQ method. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): onnx model. data_reader (data_reader.CalibrationDataReader): data_reader for calibration. weight_config (dict, optional): quantization config For example, weight_config = { '(fc2, "MatMul")': { 'weight_dtype': 'int', 'weight_bits': 4, 'weight_group_size': 32, 'weight_sym': True, 'accuracy_level': 0 }. Defaults to {}. percdamp (float, optional): percentage of Hessian's diagonal values' average, which will be added to Hessian's diagonal to increase numerical stability. Defaults to 0.01. block_size (int, optional): execute GPTQ quantization per block. Defaults to 128. actorder (bool, optional): whether to sort Hessian's diagonal values to rearrange channel-wise quantization order. Defaults to False. mse (bool, optional): whether get scale and zero point with mse error. Defaults to False. perchannel (bool, optional): whether quantize weight per-channel. Defaults to True. providers (list, optional): providers to use. Defaults to ["CPUExecutionProvider"]. return_modelproto (bool, optionmal): whether to return onnx.Modelproto. set False for layer-wise quant. Default to True Returns: onnx.ModelProto: quantized onnx model """ if not isinstance(model, onnx_model.ONNXModel): model = onnx_model.ONNXModel(model) base_dir = os.path.dirname(model.model_path) if model.model_path is not None else "" inputs, so = quant_utils.prepare_inputs(model, data_reader, providers) del data_reader org_output = copy.deepcopy(model.model.graph.output) model.remove_tensors_from_outputs([i.name for i in org_output]) output_names = [] for node in model.nodes(): # check op_type of node is MatMul # check op_name in quantization config # check dim 1 of input is weight tensor if ( node.op_type in ["MatMul"] and node.name in weight_config and model.get_initializer(node.input[1]) is not None ): output_names.append(node.input[0]) output_names = list(set(output_names)) model.add_tensors_to_outputs(output_names) if model.is_large_model: onnx.save_model( model.model, model.model_path + "_augment.onnx", save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) session = ( ort.InferenceSession(model.model.SerializeToString(), so, providers=providers) if not model.is_large_model else ort.InferenceSession(model.model_path + "_augment.onnx", so, providers=providers) ) input_name_to_nodes = model.input_name_to_nodes() for idx, input_name in enumerate(output_names): utility.simple_progress_bar(len(output_names), idx + 1) node_list = [] weights = [] for node in input_name_to_nodes[input_name]: # check op_type of node is MatMul # check op_name in quantization config # check dim 1 of input is weight tensor if ( node.op_type in ["MatMul"] and node.name in weight_config and model.get_initializer(node.input[1]) is not None ): weight = onnx.numpy_helper.to_array( model.get_initializer(model.get_node(node.name).input[1]), base_dir ).copy() if len(weight.shape) != 2: # pragma: no cover continue weights.append(weight) node_list.append(model.get_node(node.name)) if len(weights) == 0: # pragma: no cover continue Hs = [np.zeros((i.shape[0], i.shape[0])) for i in weights] nsamples = 0 for data in inputs: inp = session.run([input_name], data)[0] tmp = inp.shape[0] inp = np.reshape(inp, (-1, inp.shape[-1])) Hs = [i * (nsamples / (nsamples + tmp)) for i in Hs] nsamples += tmp inp = np.sqrt(2 / nsamples) * inp Hs = [i + np.matmul(inp.T, inp) for i in Hs] for ( node, weight, H, ) in zip(node_list, weights, Hs): num_bits = weight_config[node.name].get("weight_bits", 4) group_size = weight_config[node.name].get("weight_group_size", 32) sym = weight_config[node.name].get("weight_sym", True) dtype = weight_config[node.name].get("weight_dtype", "int") accuracy_level = weight_config[node.name].get("accuracy_level", 0) quant_format = getattr(weight_config[node.name].get("quant_format", None), "value", None) weight_tensor = model.get_initializer(node.input[1]) init_share_num = model.get_initializer_share_num(node.input[1]) # weight -> quant -> dequant -> q_weight q_weight = _gptq( weight, H, num_bits=num_bits, group_size=group_size, sym=sym, block_size=block_size, percdamp=percdamp, actorder=actorder, mse=mse, perchannel=perchannel, ) new_nodes, new_inits, remove_nodes = quant_utils.quant_matmul_weight_only( node=node, weight=weight, dtype=dtype, num_bits=num_bits, sym=sym, group_size=group_size, quant_format=quant_format, accuracy_level=accuracy_level, ) model.add_initializers(new_inits) model.add_nodes(new_nodes) model.remove_nodes(remove_nodes) if init_share_num == 1: model.remove_initializer(weight_tensor) model.remove_tensors_from_outputs(output_names) model.model.graph.output.MergeFrom(org_output) model.topological_sort() # reload external data to prevent external data file path errors if model.is_large_model: onnx.external_data_helper.load_external_data_for_model(model.model, os.path.split(model.model_path)[0]) if return_modelproto: return model.model else: model.save(model.model_path + "_quant.onnx") return model def apply_gptq_on_model( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], quant_config: dict, calibration_data_reader: data_reader.CalibrationDataReader, percdamp: float = 0.01, block_size: int = 128, actorder: bool = False, mse: bool = False, perchannel: bool = True, providers: List[str] = ["CPUExecutionProvider"], layer_wise_quant: bool = False, ) -> onnx.ModelProto: """Apply GPTQ on onnx model. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): onnx model. quant_config (dict): quantization config. calibration_data_reader (data_reader.CalibrationDataReader): data_reader for calibration. Returns: onnx.ModelProto: quantized onnx model. """ # set other model params quant_kwargs = { "percdamp": percdamp, "block_size": block_size, "actorder": actorder, "mse": mse, "perchannel": perchannel, "providers": providers, } if layer_wise_quant: quantized_model = core.layer_wise_quant( model, quant_func=gptq_quantize, weight_config=quant_config, data_reader=calibration_data_reader, **quant_kwargs ) else: quantized_model = gptq_quantize( model, data_reader=calibration_data_reader, weight_config=quant_config, **quant_kwargs ) if isinstance(quantized_model, onnx_model.ONNXModel): quantized_model = quantized_model.model quant_utils.dump_woq_stats(quantized_model, quant_config) return quantized_model ================================================ FILE: onnx_neural_compressor/algorithms/weight_only/rtn.py ================================================ # Copyright (c) 2023 MIT HAN Lab # This source code is licensed under the MIT license # # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import pathlib import numpy as np import onnx import onnxruntime as ort from onnx_neural_compressor import constants, onnx_model, utility from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor.algorithms.layer_wise import core from typing import List, Union # isort: skip def rtn_quantize( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], weight_config: dict = {}, ratios: dict = {}, providers: List[str] = ["CPUExecutionProvider"], return_modelproto: bool = True, ): """Quantize the model with round to nearst method. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): onnx model weight_config (dict, optional): quantization config For example, weight_config = { '(fc2, "MatMul")': { 'weight_dtype': 'int', 'weight_bits': 4, 'weight_group_size': 32, 'weight_sym': True, 'accuracy_level': 0 } }. Defaults to {}. ratios (dict, optional): percentile of clip. Defaults to {}. providers (list, optional): providers to use. Defaults to ["CPUExecutionProvider"]. return_modelproto (bool, optionmal): whether to return onnx.Modelproto. set False for layer-wise quant. Default to True Returns: onnx.ModelProto: quantized onnx model. """ if not isinstance(model, onnx_model.ONNXModel): model = onnx_model.ONNXModel(model) base_dir = os.path.dirname(model.model_path) if model.model_path is not None else "" new_nodes_all = [] remove_nodes_all = [] total_num = len([i for i in model.nodes() if i.op_type in ["MatMul"]]) curr_id = 0 for node in model.nodes(): if node.op_type in ["MatMul"]: curr_id += 1 utility.simple_progress_bar(total_num, curr_id) # check op_type of node is MatMul # check op_name in quantization config # check dim 1 of input is weight tensor if ( node.op_type in ["MatMul"] and node.name in weight_config and model.get_initializer(node.input[1]) is not None ): weight_tensor = model.get_initializer(node.input[1]) weight = onnx.numpy_helper.to_array(weight_tensor, base_dir=base_dir).copy() if len(weight.shape) != 2: continue dtype = weight_config[node.name].get("weight_dtype", "int") num_bits = weight_config[node.name].get("weight_bits", 4) group_size = weight_config[node.name].get("weight_group_size", 32) sym = weight_config[node.name].get("weight_sym", True) accuracy_level = weight_config[node.name].get("accuracy_level", 0) quant_format = getattr(weight_config[node.name].get("quant_format", None), "value", None) init_share_num = model.get_initializer_share_num(node.input[1]) new_nodes, new_inits, remove_nodes = quant_utils.quant_matmul_weight_only( node=node, weight=weight, dtype=dtype, num_bits=num_bits, sym=sym, group_size=group_size, ratio=ratios.get(node.input[1], 1), quant_format=quant_format, accuracy_level=accuracy_level, ) model.add_initializers(new_inits) new_nodes_all.extend(new_nodes) remove_nodes_all.extend(remove_nodes) if init_share_num == 1: model.remove_initializer(weight_tensor) model.add_nodes(new_nodes_all) model.remove_nodes(remove_nodes_all) model.topological_sort() # reload external data to prevent external data file path errors if model.is_large_model: onnx.external_data_helper.load_external_data_for_model(model.model, os.path.split(model.model_path)[0]) if return_modelproto: return model.model else: model.save(model.model_path + "_quant.onnx") return model def apply_rtn_on_model( model: Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str], quant_config: dict, ratios: dict = {}, providers: List[str] = ["CPUExecutionProvider"], layer_wise_quant: bool = False, ) -> onnx.ModelProto: """Apply RTN on onnx model. Args: model (Union[onnx.ModelProto, onnx_model.ONNXModel, pathlib.Path, str]): onnx model. quant_config (dict): quantization config. Returns: onnx.ModelProto: quantized onnx model. """ quant_kwargs = { "ratios": ratios, "providers": providers, } if layer_wise_quant: quantized_model = core.layer_wise_quant( model, quant_func=rtn_quantize, weight_config=quant_config, **quant_kwargs ) else: quantized_model = rtn_quantize(model, weight_config=quant_config, **quant_kwargs) if isinstance(quantized_model, onnx_model.ONNXModel): quantized_model = quantized_model.model quant_utils.dump_woq_stats(quantized_model, quant_config) return quantized_model ================================================ FILE: onnx_neural_compressor/quantization/__init__.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from onnx_neural_compressor.quantization.quant_utils import CalibrationMethod, QuantFormat, QuantType from onnx_neural_compressor.quantization.quantize import quantize ================================================ FILE: onnx_neural_compressor/quantization/algorithm_entry.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import pathlib import tempfile from typing import Union import onnx import onnxruntime as ort from packaging import version from onnx_neural_compressor import constants, data_reader, logger, utility from onnx_neural_compressor.algorithms.post_training_quant import calibrate, quantizer from onnx_neural_compressor.algorithms.smoother import core from onnx_neural_compressor.algorithms.weight_only import awq, gptq, rtn from onnx_neural_compressor.quantization import QuantFormat, config ort_version = version.Version(ort.__version__) ###################### RTN Algo Entry ################################## @utility.register_algo(name=constants.RTN) def rtn_quantize_entry( model: Union[pathlib.Path, str], quant_config: config.RTNConfig, *args, **kwargs ) -> onnx.ModelProto: """The main entry to apply rtn quantization.""" config_mapping = quant_config.to_config_mapping(model=model) quant_kwargs = dict( zip( quant_config.model_params_list, [getattr(quant_config, key, None) for key in quant_config.model_params_list], ) ) model = rtn.apply_rtn_on_model(model, config_mapping, **quant_kwargs) return model ###################### GPTQ Algo Entry ################################## @utility.register_algo(name=constants.GPTQ) def gptq_quantize_entry( model: Union[pathlib.Path, str], quant_config: config.GPTQConfig, calibration_data_reader: data_reader.CalibrationDataReader, *args, **kwargs, ) -> onnx.ModelProto: """The main entry to apply gptq quantization.""" assert calibration_data_reader is not None, "Please provide calibration_data_reader" assert isinstance( calibration_data_reader, data_reader.CalibrationDataReader ), "Please follow onnx_neural_compressor/data_reader.py to implement calibration_data_reader" config_mapping = quant_config.to_config_mapping(model=model) quant_kwargs = dict( zip( quant_config.model_params_list, [getattr(quant_config, key, None) for key in quant_config.model_params_list], ) ) # regenerate to ensure data exists calibration_data_reader.rewind() model = gptq.apply_gptq_on_model(model, config_mapping, calibration_data_reader, **quant_kwargs) return model ###################### AWQ Algo Entry ################################## @utility.register_algo(name=constants.AWQ) def awq_quantize_entry( model: Union[pathlib.Path, str], quant_config: config.AWQConfig, calibration_data_reader: data_reader.CalibrationDataReader, *args, **kwargs, ) -> onnx.ModelProto: """The main entry to apply awq quantization.""" assert calibration_data_reader is not None, "Please provide calibration_data_reader" assert isinstance( calibration_data_reader, data_reader.CalibrationDataReader ), "Please follow onnx_neural_compressor/data_reader.py to implement calibration_data_reader" config_mapping = quant_config.to_config_mapping(model=model) quant_kwargs = dict( zip( quant_config.model_params_list, [getattr(quant_config, key, None) for key in quant_config.model_params_list], ) ) # regenerate to ensure data exists calibration_data_reader.rewind() model = awq.apply_awq_on_model(model, config_mapping, calibration_data_reader, **quant_kwargs) return model ###################### Static quant Entry ################################## @utility.register_algo(name=constants.STATIC_QUANT) def static_quantize_entry( model: Union[pathlib.Path, str], quant_config: config.StaticQuantConfig, calibration_data_reader: data_reader.CalibrationDataReader, model_output: Union[pathlib.Path, str] = None, *args, **kwargs, ) -> onnx.ModelProto: """The main entry to apply dynamic quantization.""" if len(quant_config.op_types_to_quantize) == 0: logger.warning("No candidate op type to do quantization, exit.") exit(0) assert calibration_data_reader is not None, "Please provide calibration_data_reader" assert isinstance( calibration_data_reader, data_reader.CalibrationDataReader ), "Please follow onnx_neural_compressor/quantization/calibrate.py to implement calibration_data_reader" config_mapping = quant_config.to_config_mapping(model=model) calibration_data_reader.rewind() augment = calibrate.ONNXRTAugment( model, calibration_data_reader, dump_op_types=quant_config.op_types_to_quantize, execution_provider=quant_config.execution_provider, iterations=list(range(0, quant_config.calibration_sampling_size)), ) min_max = augment.dump_minmax(config_mapping) quantize_params = augment.dump_calibration(config_mapping, min_max=min_max) _quantizer = quantizer.StaticQuantizer( model, config_mapping, quant_format=quant_config.quant_format.name.lower(), quantization_params=quantize_params, op_types_to_quantize=quant_config.op_types_to_quantize, execution_provider=quant_config.execution_provider, optypes_to_exclude_output_quant=quant_config.optypes_to_exclude_output_quant, dedicated_qdq_pair=quant_config.dedicated_qdq_pair, add_qdq_pair_to_weight=quant_config.add_qdq_pair_to_weight, ) _quantizer.quantize_model() if model_output is not None: _quantizer.model.save(model_output) return _quantizer.model.model ###################### SmoothQuant Entry ################################## @utility.register_algo(name=constants.SMOOTH_QUANT) def smooth_quant_entry( model: Union[pathlib.Path, str], quant_config: config.SmoothQuantConfig, calibration_data_reader: data_reader.CalibrationDataReader, model_output: Union[pathlib.Path, str] = None, *args, **kwargs, ) -> Union[pathlib.Path, str, onnx.ModelProto]: """Apply smooth quant.""" assert calibration_data_reader is not None, "Please provide calibration_data_reader" assert isinstance( calibration_data_reader, data_reader.CalibrationDataReader ), "Please follow onnx_neural_compressor/data_reader.py to implement calibration_data_reader" # smooth operation calibration_data_reader.rewind() smoother = core.Smoother( model, calibration_data_reader, execution_provider=getattr(quant_config, "execution_provider", "CPUExecutionProvider"), ) smoothed_model = smoother.transform(**quant_config.get_model_params_dict()) with tempfile.TemporaryDirectory(prefix="ort.quant.") as tmp_dir: # ORT quant API requires str input onnx.save_model( smoothed_model, pathlib.Path(tmp_dir).joinpath("smooth.onnx").as_posix(), save_as_external_data=True, all_tensors_to_one_file=True, location="smooth.onnx_data", size_threshold=1024, convert_attribute=False, ) # quant operation calibration_data_reader.rewind() # exclude Mul operations which are inserted during smooth operation excluded_nodes = [i.name for i in smoothed_model.graph.node if i.name.endswith("_smooth_mul")] quant_config.nodes_to_exclude.extend(excluded_nodes) q_model = static_quantize_entry( pathlib.Path(tmp_dir).joinpath("smooth.onnx").as_posix(), quant_config, calibration_data_reader, model_output, ) return q_model ###################### Dynamic quant Entry ################################## @utility.register_algo(name=constants.DYNAMIC_QUANT) def dynamic_quantize_entry( model: Union[pathlib.Path, str], quant_config: config.DynamicQuantConfig, model_output: Union[pathlib.Path, str] = None, *args, **kwargs, ) -> onnx.ModelProto: """The main entry to apply dynamic quantization.""" if len(quant_config.op_types_to_quantize) == 0: logger.warning("No candidate op type to do quantization, exit.") exit(0) config_mapping = quant_config.to_config_mapping(model=model) _quantizer = quantizer.DynamicQuantizer( model, config_mapping, op_types_to_quantize=quant_config.op_types_to_quantize, ) _quantizer.quantize_model() if model_output is not None: _quantizer.model.save(model_output) return _quantizer.model.model ================================================ FILE: onnx_neural_compressor/quantization/matmul_4bits_quantizer.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import List, Union # isort: skip import onnx import onnxruntime as ort from onnx_neural_compressor.quantization import matmul_nbits_quantizer RTNWeightOnlyQuantConfig = matmul_nbits_quantizer.RTNWeightOnlyQuantConfig AWQWeightOnlyQuantConfig = matmul_nbits_quantizer.AWQWeightOnlyQuantConfig GPTQWeightOnlyQuantConfig = matmul_nbits_quantizer.GPTQWeightOnlyQuantConfig class MatMul4BitsQuantizer(matmul_nbits_quantizer.MatMulNBitsQuantizer): def __init__( self, model: Union[onnx.ModelProto, str], block_size: int = 128, is_symmetric: bool = False, is_signed: bool = False, accuracy_level: int = 0, nodes_to_exclude=None, algo_config: matmul_nbits_quantizer.WeightOnlyQuantConfig = None, providers: List[str] = ["CPUExecutionProvider"], optimization_level: ort.GraphOptimizationLevel = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC, ): super().__init__( model=model, block_size=block_size, is_symmetric=is_symmetric, is_signed=is_signed, accuracy_level=accuracy_level, nodes_to_exclude=nodes_to_exclude, algo_config=algo_config, n_bits=4, providers=providers, optimization_level=optimization_level, ) ================================================ FILE: onnx_neural_compressor/quantization/matmul_nbits_quantizer.py ================================================ # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import List, Union # isort: skip import pathlib import tempfile import onnx import onnxruntime as ort from onnx_neural_compressor import data_reader, logger, onnx_model, utility from onnx_neural_compressor.quantization import QuantFormat from onnx_neural_compressor.quantization import algorithm_entry as algos from onnx_neural_compressor.quantization import config class WeightOnlyQuantConfig: def __init__(self, algorithm, quant_format=QuantFormat.QOperator): """This is the Base class for Weight Only Quant Configuration. Args: algorithm: weight only quantize algorithm name. """ self.algorithm = algorithm self.quant_format = quant_format class RTNWeightOnlyQuantConfig(WeightOnlyQuantConfig): def __init__(self, ratios=None, layer_wise_quant=False, quant_format=QuantFormat.QOperator): super().__init__( algorithm="RTN", quant_format=quant_format, ) if ratios is None: ratios = {} self.ratios = ratios self.layer_wise_quant = layer_wise_quant class GPTQWeightOnlyQuantConfig(WeightOnlyQuantConfig): def __init__( self, calibration_data_reader: data_reader.CalibrationDataReader, percdamp=0.01, block_size=128, actorder=False, mse=False, perchannel=True, layer_wise_quant=False, quant_format=QuantFormat.QOperator, ): super().__init__( algorithm="GPTQ", quant_format=quant_format, ) self.calibration_data_reader = calibration_data_reader self.percdamp = percdamp self.block_size = block_size self.actorder = actorder self.mse = mse self.perchannel = perchannel self.layer_wise_quant = layer_wise_quant class AWQWeightOnlyQuantConfig(WeightOnlyQuantConfig): def __init__( self, calibration_data_reader: data_reader.CalibrationDataReader, enable_auto_scale=True, enable_mse_search=True, quant_format=QuantFormat.QOperator, ): super().__init__(algorithm="AWQ", quant_format=quant_format) self.calibration_data_reader = calibration_data_reader self.enable_auto_scale = enable_auto_scale self.enable_mse_search = enable_mse_search algorithm_config_mapping = { "RTN": RTNWeightOnlyQuantConfig, "AWQ": AWQWeightOnlyQuantConfig, "GPTQ": GPTQWeightOnlyQuantConfig, } class MatMulNBitsQuantizer: def __init__( self, model: Union[onnx.ModelProto, str], block_size: int = 128, is_symmetric: bool = False, is_signed: bool = False, accuracy_level: int = 0, nodes_to_exclude: List[str] = None, algo_config: WeightOnlyQuantConfig = None, n_bits: int = 4, providers: List[str] = ["CPUExecutionProvider"], optimization_level: ort.GraphOptimizationLevel = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC, ): if nodes_to_exclude is None: nodes_to_exclude = [] self.model = model self.block_size = block_size self.is_symmetric = is_symmetric self.is_signed = is_signed self.accuracy_level = accuracy_level self.nodes_to_exclude = list(set(nodes_to_exclude)) self.algo_config = algo_config or RTNWeightOnlyQuantConfig() self.n_bits = n_bits self.providers = providers self.algorithm = self.algo_config.algorithm self.optimization_level = optimization_level assert self.algorithm in [ "RTN", "AWQ", "GPTQ", ], "Only RTN, GPTQ and AWQ algorithms are supported, but get {} algorithm".format(self.algorithm) def _generate_nc_config(self): config_class = config.config_registry.get_cls_configs()[self.algorithm.lower()] quant_kwargs = { "weight_dtype": "int" if self.is_signed else "uint", "weight_bits": self.n_bits, "weight_group_size": self.block_size, "weight_sym": self.is_symmetric, "accuracy_level": self.accuracy_level, "providers": self.providers, "quant_format": self.algo_config.quant_format, "nodes_to_exclude": self.nodes_to_exclude, } if self.algorithm == "RTN": quant_kwargs.update( { "layer_wise_quant": self.algo_config.layer_wise_quant, } ) elif self.algorithm == "GPTQ": quant_kwargs.update( { "percdamp": self.algo_config.percdamp, "block_size": self.algo_config.block_size, "actorder": self.algo_config.actorder, "mse": self.algo_config.mse, "perchannel": self.algo_config.perchannel, "layer_wise_quant": self.algo_config.layer_wise_quant, } ) elif self.algorithm == "AWQ": quant_kwargs.update( { "enable_auto_scale": self.algo_config.enable_auto_scale, "enable_mse_search": self.algo_config.enable_mse_search, } ) nc_config = config_class(**quant_kwargs) return nc_config def int4_quant_algo(self): qconfig = self._generate_nc_config() model = self.model opt_tmp_file = tempfile.TemporaryDirectory() if getattr(self.algo_config, "layer_wise_quant", False) and not isinstance(model, str): logger.warning("Please use model path for layer-wise quantization.") # do graph optimization if not layer_wise_quant if ( not getattr(self.algo_config, "layer_wise_quant", False) and self.optimization_level != ort.GraphOptimizationLevel.ORT_DISABLE_ALL ): if not isinstance(model, str): onnx.save_model( model, pathlib.Path(opt_tmp_file.name).joinpath("tmp.onnx").as_posix(), save_as_external_data=True, all_tensors_to_one_file=True, location="tmp.onnx_data", size_threshold=1024, convert_attribute=False, ) model = pathlib.Path(opt_tmp_file.name).joinpath("tmp.onnx").as_posix() logger.info("Start graph optimization...") sess_options = ort.SessionOptions() sess_options.graph_optimization_level = self.optimization_level sess_options.optimized_model_filepath = pathlib.Path(opt_tmp_file.name).joinpath("opt.onnx").as_posix() sess_options.add_session_config_entry( "session.optimized_model_external_initializers_file_name", "opt.onnx_data" ) sess_options.add_session_config_entry( "session.optimized_model_external_initializers_min_size_in_bytes", "1024" ) session = ort.InferenceSession(model, sess_options, providers=["CPUExecutionProvider"]) model = sess_options.optimized_model_filepath del session logger.info("Graph optimization done.") logger.info(f"start to quantize model with {self.algorithm} algorithm...") if self.algorithm == "RTN": self.model = algos.rtn_quantize_entry(model, qconfig) elif self.algorithm == "GPTQ": self.model = algos.gptq_quantize_entry(model, qconfig, self.algo_config.calibration_data_reader) elif self.algorithm == "AWQ": self.model = algos.awq_quantize_entry(model, qconfig, self.algo_config.calibration_data_reader) logger.info(f"complete quantization of model with {self.algorithm} algorithm.") opt_tmp_file.cleanup() def process(self): self.int4_quant_algo() ================================================ FILE: onnx_neural_compressor/quantization/quant_utils.py ================================================ # Copyright (c) 2023 MIT HAN Lab # This source code is licensed under the MIT license # # Copyright (c) 2024 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import enum import onnx class QuantType(enum.Enum): # pragma: no cover """Represent QuantType value.""" QInt8 = 0 QUInt8 = 1 QInt4 = 4 QUInt4 = 5 @property def tensor_type(self): if self == QuantType.QInt8: return onnx.TensorProto.INT8 if self == QuantType.QUInt8: return onnx.TensorProto.UINT8 if self == QuantType.QInt8: return onnx.TensorProto.INT4 if self == QuantType.QUInt4: return onnx.TensorProto.UINT4 raise ValueError(f"Unexpected value qtype={self!r}.") class QuantFormat(enum.Enum): QOperator = 0 QDQ = 1 class CalibrationMethod(enum.Enum): MinMax = 0 Entropy = 1 Percentile = 2 Distribution = 3 ================================================ FILE: onnx_neural_compressor/quantization/quantize.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import pathlib import tempfile from typing import Union import onnx import onnxruntime as ort from onnxruntime.quantization.quantize import QuantConfig from onnx_neural_compressor.quantization import algorithm_entry as algos from onnx_neural_compressor.quantization import config # ORT-like user-facing API def quantize( model_input: Union[str, pathlib.Path, onnx.ModelProto], model_output: Union[str, pathlib.Path], quant_config: config.BaseConfig, optimization_level: ort.GraphOptimizationLevel = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC, ): with tempfile.TemporaryDirectory(prefix="ort.opt.") as tmp_dir: if optimization_level != ort.GraphOptimizationLevel.ORT_DISABLE_ALL: sess_options = ort.SessionOptions() sess_options.graph_optimization_level = optimization_level sess_options.optimized_model_filepath = pathlib.Path(tmp_dir).joinpath("opt.onnx").as_posix() sess_options.add_session_config_entry( "session.optimized_model_external_initializers_file_name", "opt.onnx_data" ) sess_options.add_session_config_entry( "session.optimized_model_external_initializers_min_size_in_bytes", "1024" ) session = ort.InferenceSession(model_input, sess_options, provides=["CPUExecutionProvider"]) del session model_input = sess_options.optimized_model_filepath if isinstance(quant_config, config.StaticQuantConfig): if quant_config.extra_options.get("SmoothQuant", False): algos.smooth_quant_entry( model_input, quant_config, quant_config.calibration_data_reader, model_output=model_output ) else: algos.static_quantize_entry( model_input, quant_config, quant_config.calibration_data_reader, model_output=model_output ) elif isinstance(quant_config, config.DynamicQuantConfig): algos.dynamic_quantize_entry(model_input, quant_config, model_output=model_output) else: raise TypeError( "Invalid quantization config type, it must be either StaticQuantConfig or DynamicQuantConfig." ) ================================================ FILE: onnx_neural_compressor/quantization/tuning.py ================================================ # Copyright (c) 2023 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import os import pathlib import shutil import tempfile import traceback import uuid import onnx import onnxruntime as ort from onnx import external_data_helper from onnx_neural_compressor import data_reader, logger, utility from onnx_neural_compressor.quantization import config from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, Sized, Tuple, Union # isort: skip class EvaluationFuncWrapper: def __init__(self, eval_fn: Callable, eval_args=None): """Evaluation function wrapper. Args: eval_fn: a function for evaluated the float or quantized model eval_args: positional arguments for `eval_fn` """ self.eval_fn = eval_fn self.eval_args = eval_args def evaluate(self, model) -> Union[float, int]: result = self.eval_fn(model, *self.eval_args) if self.eval_args else self.eval_fn(model) return result class Evaluator: """Evaluator is a collection of evaluation functions. Note: will deprecate this class in the future. Examples: def eval_acc(model): ... def eval_perf(molde): ... # Usage user_eval_fns1 = eval_acc user_eval_fns2 = {"eval_fn": eval_acc} user_eval_fns3 = {"eval_fn": eval_acc, "weight": 1.0, "name": "accuracy"} user_eval_fns4 = [ {"eval_fn": eval_acc, "weight": 0.5}, {"eval_fn": eval_perf, "weight": 0.5, "name": "accuracy"}, ] """ EVAL_FN = "eval_fn" WEIGHT = "weight" FN_NAME = "name" EVAL_FN_TEMPLATE: Dict[str, Any] = {EVAL_FN: None, WEIGHT: 1.0, FN_NAME: None} def __init__(self) -> None: self.eval_fn_registry: List[Dict[str, Any]] = [] def evaluate(self, model) -> float: """Evaluate the model using registered evaluation functions. Args: model: The fp32 model or quantized model. Returns: The overall result of all registered evaluation functions. """ result = 0 for eval_pair in self.eval_fn_registry: eval_fn = eval_pair[self.EVAL_FN] eval_result = eval_fn(model) result = self._update_the_objective_score(eval_pair, eval_result, result) return result def _update_the_objective_score(self, eval_pair, eval_result, overall_result) -> float: return overall_result + eval_result * eval_pair[self.WEIGHT] def get_number_of_eval_functions(self) -> int: return len(self.eval_fn_registry) def _set_eval_fn_registry(self, user_eval_fns: List[Dict]) -> None: self.eval_fn_registry = [ { self.EVAL_FN: user_eval_fn_pair[self.EVAL_FN], self.WEIGHT: user_eval_fn_pair.get(self.WEIGHT, 1.0), self.FN_NAME: user_eval_fn_pair.get( self.FN_NAME, getattr(user_eval_fn_pair[self.EVAL_FN], "__name__", "custom_func") ), } for user_eval_fn_pair in user_eval_fns ] def set_eval_fn_registry(self, eval_fns: Optional[Union[Callable, Dict, List[Dict]]] = None) -> None: # About the eval_fns format, refer the class docstring for details. if eval_fns is None: return elif callable(eval_fns): # single eval_fn eval_fn_pair = copy.deepcopy(self.EVAL_FN_TEMPLATE) eval_fn_pair[self.EVAL_FN] = eval_fns eval_fn_pair[self.FN_NAME] = eval_fns.__name__ eval_fns = [eval_fn_pair] elif isinstance(eval_fns, Dict): eval_fns = [eval_fns] elif isinstance(eval_fns, List): assert all([isinstance(eval_fn_pair, Dict) for eval_fn_pair in eval_fns]) else: raise NotImplementedError(f"The eval_fns should be a dict or a list of dict, but got {type(eval_fns)}.") self._set_eval_fn_registry(eval_fns) def self_check(self) -> None: # check the number of evaluation functions num_eval_fns = self.get_number_of_eval_functions() assert num_eval_fns > 0, "Please ensure that you register at least one evaluation metric for auto-tune." logger.info("There are %d evaluations functions.", num_eval_fns) evaluator = Evaluator() class ConfigSet: def __init__(self, config_list: List[config.BaseConfig]) -> None: self.config_list = config_list def __getitem__(self, index) -> config.BaseConfig: assert 0 <= index < len(self.config_list), f"Index {index} out of range." return self.config_list[index] def __len__(self) -> int: return len(self.config_list) @classmethod def _from_single_config(cls, fwk_config: config.BaseConfig) -> List[config.BaseConfig]: config_list = [] config_list = fwk_config.expand() return config_list @classmethod def _from_list_of_configs(cls, fwk_configs: List[config.BaseConfig]) -> List[config.BaseConfig]: config_list = [] for fwk_config in fwk_configs: config_list += cls._from_single_config(fwk_config) return config_list @classmethod def generate_config_list(cls, fwk_configs: Union[config.BaseConfig, List[config.BaseConfig]]): # There are several cases for the input `fwk_configs`: # 1. fwk_configs is a single config # 2. fwk_configs is a list of configs # For a single config, we need to check if it can be expanded or not. config_list = [] if isinstance(fwk_configs, config.BaseConfig): config_list = cls._from_single_config(fwk_configs) elif isinstance(fwk_configs, List): config_list = cls._from_list_of_configs(fwk_configs) else: raise NotImplementedError(f"Unsupported type {type(fwk_configs)} for fwk_configs.") return config_list @classmethod def from_fwk_configs(cls, fwk_configs: Union[config.BaseConfig, List[config.BaseConfig]]) -> "ConfigSet": """Create a ConfigSet object from a single config or a list of configs. Args: fwk_configs: A single config or a list of configs. Examples: 1) single config: config.RTNConfig(weight_group_size=32) 2) single expandable config: config.RTNConfig(weight_group_size=[32, 64]) 3) mixed 1) and 2): [config.RTNConfig(weight_group_size=32), config.RTNConfig(weight_group_size=[32, 64])] Returns: ConfigSet: A ConfigSet object. """ config_list = cls.generate_config_list(fwk_configs) return cls(config_list) class Sampler: def __init__(self, config_source: Optional[ConfigSet]) -> None: pass def __iter__(self) -> Iterator[config.BaseConfig]: """Iterate over indices of config set elements.""" raise NotImplementedError class SequentialSampler(Sampler): """Samples elements sequentially, always in the same order. Args: config_source (_ConfigSet): config set to sample from """ config_source: Sized def __init__(self, config_source: Sized) -> None: self.config_source = config_source def __iter__(self) -> Iterator[int]: return iter(range(len(self.config_source))) def __len__(self) -> int: return len(self.config_source) default_sampler = SequentialSampler class ConfigLoader: def __init__( self, config_set: ConfigSet, sampler: Sampler = default_sampler, skip_verified_config: bool = True ) -> None: self.config_set = ConfigSet.from_fwk_configs(config_set) self._sampler = sampler(self.config_set) self.skip_verified_config = skip_verified_config self.verify_config_list = list() def is_verified_config(self, config): for verified_config in self.verify_config_list: if config == verified_config: return True return False def __iter__(self) -> Generator[config.BaseConfig, Any, None]: for index in self._sampler: new_config = self.config_set[index] if self.skip_verified_config and self.is_verified_config(new_config): logger.debug("Skip the verified config:") logger.debug(new_config.to_dict()) continue self.verify_config_list.append(new_config) yield new_config class TuningConfig: """Config for auto tuning pipeline. Examples: from onnx_neural_compressor.quantization import tuning tune_config = tuning.TuningConfig( config_set=[config1, config2, ...], max_trials=3, tolerable_loss=0.01) The tuning process stops when either of the following conditions is met: 1) The number of trials reaches the maximum trials. 2) The metric loss is within the tolerable loss. For condition 2), we calculate the metric loss as follows: relative_loss = (fp32_baseline - eval_result_of_q_model) / fp32_baseline If relative_loss <= tolerable_loss, we stop the tuning process. For example: tolerable_loss = 0.01 fp32_baseline = 100 eval_result_of_q_model = 99 relative_loss = (100 - 99) / 100 = 0.01 The metric loss is within the tolerable loss, so the tuning process is stopped. """ def __init__( self, config_set: Union[config.BaseConfig, List[config.BaseConfig]] = None, sampler: Sampler = default_sampler, tolerable_loss=0.01, max_trials=100, ): """Initial a TuningConfig. Args: config_set: A single config or a list of configs. Defaults to None. sampler: tuning sampler that decide the trials order. Defaults to default_sampler. tolerable_loss: This float indicates how much metric loss we can accept. The metric loss is relative, it can be both positive and negative. Default is 0.01. max_trials: Max tuning times. Combine with `tolerable_loss` field to decide when to stop. Default is 100. """ self.config_set = config_set self.sampler = sampler self.tolerable_loss = tolerable_loss self.max_trials = max_trials class _TrialRecord: @staticmethod def _generate_unique_id(): unique_id = str(uuid.uuid4()) return unique_id def __init__(self, trial_index: int, trial_result: Union[int, float], quant_config: config.BaseConfig): # The unique id to refer to one trial self.trial_id = _TrialRecord._generate_unique_id() self.trial_index = trial_index self.trial_result = trial_result self.quant_config = quant_config class TuningMonitor: def __init__(self, tuning_config: TuningConfig) -> None: self.tuning_config = tuning_config self.trial_cnt = 0 self.tuning_history: List[_TrialRecord] = [] self.baseline = None def add_trial_result( self, trial_index: int, trial_result: Union[int, float], quant_config: config.BaseConfig ) -> None: self.trial_cnt += 1 trial_record = _TrialRecord(trial_index, trial_result, quant_config) self.tuning_history.append(trial_record) def set_baseline(self, baseline: float): self.baseline = baseline logger.info(f"Fp32 baseline is {self.baseline}") def get_number_of_trials(self): return len(self.tuning_history) def need_skip(self, config) -> bool: """Check whether the expanded quant config is verified.""" if len(self.tuning_history) > 0 and any([config == i.quant_config.config_mapping for i in self.tuning_history]): logger.warning("Skip the verified config mapping.") logger.debug(config) return True return False def need_stop(self) -> bool: """Check if need to stop tuning. Either accuracy goal is met, max trials is reached or timeout is reached. Returns: stop_flag: True if need to stop, otherwise False. """ # reach max trials reach_max_trials = self.trial_cnt >= self.tuning_config.max_trials # reach accuracy goal meet_accuracy_goal = ( False if self.baseline is None else self.tuning_history[-1].trial_result >= (self.baseline * (1 - self.tuning_config.tolerable_loss)) ) # [-1] is the last element representing the latest trail record. return reach_max_trials or meet_accuracy_goal def print_config_diff(self, config): if len(self.tuning_history) == 0: logger.info("quant config: {}".format(config)) else: logger.info("quant config difference: {}".format(config.get_diff_dict(self.tuning_history[0].quant_config))) class TuningLogger: """A unified logger for the tuning/quantization process. It assists validation teams in retrieving logs. """ @classmethod def tuning_start(cls) -> None: logger.info("Tuning started.") @classmethod def trial_start(cls, trial_index: int = None) -> None: logger.info("%d-trail started.", trial_index) @classmethod def quantization_start(cls, stacklevel=2) -> None: logger.info("Quantization started.", stacklevel=stacklevel) @classmethod def quantization_end(cls, stacklevel=2) -> None: logger.info("Quantization end.", stacklevel=stacklevel) @classmethod def evaluation_start(cls) -> None: logger.info("Evaluation started.") @classmethod def evaluation_end(cls) -> None: logger.info("Evaluation end.") @classmethod def trial_end(cls, trial_index: int = None) -> None: logger.info("%d-trail end.", trial_index) @classmethod def tuning_end(cls) -> None: logger.info("Tuning completed.") def init_tuning(tuning_config: TuningConfig) -> Tuple[ConfigLoader, TuningLogger, TuningMonitor]: config_loader = ConfigLoader(config_set=tuning_config.config_set, sampler=tuning_config.sampler) tuning_logger = TuningLogger() tuning_monitor = TuningMonitor(tuning_config) return config_loader, tuning_logger, tuning_monitor def get_all_config_set() -> Union[config.BaseConfig, List[config.BaseConfig]]: return config.get_all_config_set_from_config_registry() def _need_apply(quant_config: config.BaseConfig, algo_name): return quant_config.name == algo_name if hasattr(quant_config, "name") else False def _quantize( model_input: Union[pathlib.Path, str], quant_config: config.BaseConfig, calibration_data_reader: data_reader.CalibrationDataReader = None, ) -> onnx.ModelProto: """The main entry to quantize a model. Args: model_input (Union[pathlib.Path, str]): Path or str to the model to quantize. quant_config (config.BaseConfig): a quantization configuration. calibration_data_reader (data_reader.CalibrationDataReader, optional): dataloader for calibration. Defaults to None. Returns: onnx.ModelProto: The quantized model. """ registered_configs = config.config_registry.get_cls_configs() if isinstance(quant_config, dict): quant_config = config.ComposableConfig.from_dict(quant_config, config_registry=registered_configs) logger.info(f"Parsed a config dict to construct the quantization config: {quant_config}.") else: assert isinstance( quant_config, config.BaseConfig ), f"Please pass a dict or config instance as the quantization configuration, but got {type(quant_config)}." logger.debug(f"Quantize model with config: \n {quant_config} \n") # select quantization algo according to config q_model = None for algo_name, algo_func in utility.algos_mapping.items(): if _need_apply(quant_config, algo_name): logger.info(f"Start to apply {algo_name} on the model.") q_model = algo_func(model_input, quant_config, calibration_data_reader=calibration_data_reader) return q_model def autotune( model_input: Union[pathlib.Path, str], tune_config: TuningConfig, eval_fn: Callable, eval_args: Optional[Tuple[Any]] = None, calibration_data_reader: data_reader.CalibrationDataReader = None, optimization_level: ort.GraphOptimizationLevel = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC, ) -> Union[None, onnx.ModelProto]: """The main entry of auto-tune. Args: model_input (Union[pathlib.Path, str]): onnx model path. tune_config (TuningConfig): tuning config. TuningConfig is created with algorithm configs, parameters supported tuning are in their params_list. Support: Expand parameters to a list of parameters like TuningConfig(config_set=[config.RTNConfig(weight_bits=[4, 8])]) Pass a list of configs like TuningConfig(config_set=[config.RTNConfig(), config.GPTQConfig()]) eval_fn (Callable): evaluate function. During evaluation, autotune will only pass model path as the input of function. eval_args (Optional[Tuple[Any]]): evaluate arguments. Positional arguments for `eval_fn`. calibration_data_reader (data_reader.CalibrationDataReader): dataloader for calibration. optimization_level (onnxruntime.GraphOptimizationLevel): graph optimization level. Support ORT_DISABLE_ALL, ORT_ENABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_EXTENDED. Default is ORT_ENABLE_BASIC. Details: https://onnxruntime.ai/docs/performance/model-optimizations/graph-optimizations.html#onlineoffline-mode """ best_quant_model = None eval_func_wrapper = EvaluationFuncWrapper(eval_fn, eval_args) config_loader, tuning_logger, tuning_monitor = init_tuning(tuning_config=tune_config) tmp_folder = tempfile.TemporaryDirectory() pathlib.Path(tmp_folder.name).joinpath("./eval").mkdir() if optimization_level != ort.GraphOptimizationLevel.ORT_DISABLE_ALL: sess_options = ort.SessionOptions() sess_options.graph_optimization_level = optimization_level sess_options.optimized_model_filepath = pathlib.Path(tmp_folder.name).joinpath("model.onnx").as_posix() sess_options.add_session_config_entry( "session.optimized_model_external_initializers_file_name", "model.onnx_data" ) sess_options.add_session_config_entry("session.optimized_model_external_initializers_min_size_in_bytes", "1024") session = ort.InferenceSession(model_input, sess_options, providers=["CPUExecutionProvider"]) # copy config.json to tmp dir for evaluation, LLMs evaluation may need it if isinstance(model_input, str) and os.path.exists( pathlib.Path(model_input).parent.joinpath("config.json").as_posix() ): shutil.copyfile( pathlib.Path(model_input).parent.joinpath("config.json").as_posix(), pathlib.Path(tmp_folder.name).joinpath("config.json").as_posix(), ) model_input = sess_options.optimized_model_filepath del session try: baseline: float = eval_func_wrapper.evaluate(model_input) except Exception as e: if "'str' object has no attribute 'SerializeToString'" in str(e): logger.warning("Please refine your eval_fn to accept model path (str) as input.") if "Unable to load from type ''" in str(e): logger.warning("Please pass model path to autotune API rather than onnx.ModelProto.") print(traceback.format_exc()) exit(0) tuning_monitor.set_baseline(baseline) tuning_logger.tuning_start() for trial_index, quant_config in enumerate(config_loader): # check whether config_mapping is verified config_mapping = quant_config.to_config_mapping(model=model_input) if tuning_monitor.need_skip(config_mapping): continue if calibration_data_reader is not None: calibration_data_reader.rewind() tuning_logger.trial_start(trial_index=trial_index) tuning_logger.quantization_start() tuning_monitor.print_config_diff(quant_config) q_model = _quantize(model_input, quant_config=quant_config, calibration_data_reader=calibration_data_reader) tuning_logger.quantization_end() tuning_logger.evaluation_start() # evaluate API requires str input onnx.save_model( q_model, pathlib.Path(tmp_folder.name).joinpath("./eval/model.onnx").as_posix(), save_as_external_data=True, all_tensors_to_one_file=True, size_threshold=1024, convert_attribute=False, ) # copy config.json to tmp dir for evaluation, LLMs evaluation may need it if isinstance(model_input, str) and os.path.exists( pathlib.Path(model_input).parent.joinpath("config.json").as_posix() ): shutil.copyfile( pathlib.Path(model_input).parent.joinpath("config.json").as_posix(), pathlib.Path(tmp_folder.name).joinpath("./eval/config.json").as_posix(), ) eval_result: float = eval_func_wrapper.evaluate( pathlib.Path(tmp_folder.name).joinpath("./eval/model.onnx").as_posix() ) tuning_logger.evaluation_end() logger.info("Evaluation result: %.4f", eval_result) tuning_monitor.add_trial_result(trial_index, eval_result, quant_config) tuning_logger.trial_end(trial_index) if tuning_monitor.need_stop(): external_data_helper.load_external_data_for_model( q_model, pathlib.Path(tmp_folder.name).joinpath("./eval").as_posix() ) best_quant_model = q_model break tuning_logger.tuning_end() if best_quant_model is None: logger.info( "Don't find the quantized model which meets accuracy requirement. " "Please try other configs or adjust tolerable_loss." ) exit(0) tmp_folder.cleanup() return best_quant_model