Hugging Face's logo

rogermt
/
neurogolf-solver

ONNX
ml-intern
Model card Files
xet
 Community
neurogolf-solver / onnx-neural-compressor.txt
rogermt's picture
rogermt
Upload onnx-neural-compressor.txt
ae72913 verified
Raw
History Blame
447 kB
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:
<a target="_blank" href="imgs/workflow.png">
 <img src="imgs/workflow.png" alt="Workflow">
</a>
## 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
<a target="_blank" href="imgs/architecture.png">
 <img src="imgs/architecture.png" alt="Architecture" width=914 height=350>
</a>
## Workflow
<a target="_blank" href="imgs/workflow.png">
 <img src="imgs/workflow.png" alt="Workflow" width=914 height=390>
</a>
================================================
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.
<table class="center">
 <thead>
 <tr>
 <th>Backend</th>
 <th>Backend Library</th>
 <th>Support Device(cpu as default)</th>
 </tr>
 </thead>
 <tbody>
 <tr>
 <td align="left">CPUExecutionProvider</td>
 <td align="left">MLAS</td>
 <td align="left">cpu</td>
 </tr>
 <tr>
 <td align="left">TensorrtExecutionProvider</td>
 <td align="left">TensorRT</td>
 <td align="left">gpu</td>
 </tr>
 <tr>
 <td align="left">CUDAExecutionProvider</td>
 <td align="left">CUDA</td>
 <td align="left">gpu</td>
 </tr>
 <tr>
 <td align="left">DnnlExecutionProvider</td>
 <td align="left">OneDNN</td>
 <td align="left">cpu</td>
 </tr>
 <tr>
 <td align="left">DmlExecutionProvider*</td>
 <td align="left">OneDNN</td>
 <td align="left">npu</td>
 </tr>
 </tbody>
</table>
<br>
<br>
> ***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.
<img src="./imgs/lwq_ort.png" width=900 height=400>
*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
<table class="tg">
<thead>
 <tr>
 <th colspan="2" style="text-align:center;vertical-align:middle">Types/Framework</th>
 <th style="text-align:center;vertical-align:middle">ONNX Runtime</th>
 </tr>
</thead>
<tbody>
 <tr>
 <td style="text-align:center;vertical-align:middle" colspan="2">W8A8 Post Training Static Quantization</td>
 <td style="text-align:center;vertical-align:middle">&#10005;</td>
 </tr>
 <tr>
 <td style="text-align:center;vertical-align:middle" rowspan="3">Weight-only Quantization</td>
 <td style="text-align:center;vertical-align:middle">RTN</td>
 <td style="text-align:center;vertical-align:middle">&#10004;</td></td>
 </tr>
 <tr>
 <td style="text-align:center;vertical-align:middle">AWQ</td>
 <td style="text-align:center;vertical-align:middle">&#10005;</td>
 </tr>
 <tr>
 <td style="text-align:center;vertical-align:middle">GPTQ</td>
 <td style="text-align:center;vertical-align:middle">&#10004;</td>
 </tr>
</tbody>
</table>
## 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 | &#10004; |
| AWQ | &#10004; |
| GPTQ | &#10004; |
> **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"} <br> & 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.
<div align="center">
 <img src="./imgs/sq_pc.png"/>
</div>
## 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.
<div align="center">
 <img src="./imgs/sq_convert.png"/>
</div>
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.
<div align="center">
 <img src="./imgs/smoothquant.png" height="250"/>
</div>
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 '<class 'onnx.onnx_ml_pb2.ModelProto'>'" 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