google-ai-edge/ai-edge-quantizer
GitHub: google-ai-edge/ai-edge-quantizer
一款灵活的训练后量化工具，旨在帮助开发者在边缘设备上优化和部署 LiteRT 模型。
Stars: 182 | Forks: 33
# AI Edge Quantizer
一款面向高级开发者的量化工具，用于量化已转换的 LiteRT 模型。它
旨在帮助高级用户在资源
密集型模型（例如 GenAI 模型）上追求最佳性能。
## 构建状态
构建类型 | 状态 |
----------- | --------------|
单元测试 | [](https://github.com/google-ai-edge/ai-edge-quantizer/actions/workflows/nightly_unittests.yml) |
每日发布版 | [](https://github.com/google-ai-edge/ai-edge-quantizer/actions/workflows/nightly_release.yml) |
每日 Colab | [](https://github.com/google-ai-edge/ai-edge-quantizer/actions/workflows/nightly_colabs.yml) |
## 安装
### 环境要求与依赖
* Python 版本：3.10, 3.11, 3.12, 3.13
* 操作系统：Linux, MacOS
* TensorFlow: [](https://pypi.org/project/tf-nightly/)
### 安装
Nightly PyPi 包：
```
pip install ai-edge-quantizer-nightly
```
## API 用法
量化器需要两个输入：
1. 未量化的源 LiteRT 模型（FlatBuffer 格式，FP32 数据类型，扩展名为 `.tflite`）
2. 量化配置（详见下文）
并输出一个准备好部署在边缘设备上的量化 LiteRT 模型。
### 基本用法
简而言之，量化器按以下步骤运行：
1. 实例化 `Quantizer` 类。这是用户访问量化器
功能的入口点。
2. 加载所需的量化配置（详见小节）。
3. 量化（并保存）模型。量化器的
大部分内部逻辑都在此阶段运行。
```
from ai_edge_quantizer import quantizer, recipe
qt = quantizer.Quantizer("path/to/input/tflite")
# 加载一个开箱即用的 recipe（例如 dynamic int8 quantization）。
qt.load_quantization_recipe(recipe.dynamic_wi8_afp32())
qt.quantize().export_model("/path/to/output/tflite")
```
请参阅[入门 Colab](colabs/getting_started.ipynb) 获取这 3 个步骤最简单的快速入门指南，以及[选择性量化 Colab](colabs/selective_quantization_isnet.ipynb) 了解有关高级功能的更多细节。
#### LiteRT 模型
请参阅 [LiteRT 文档](https://ai.google.dev/edge/litert) 了解如何从 Jax、PyTorch 和 TensorFlow 生成 LiteRT 模型。输入的源模型应为 FlatBuffer 格式、扩展名为 `.tflite` 的 FP32（未量化）模型。
#### 量化配置
用户需要使用 AI Edge Quantizer 的 API 指定一个量化配置，以
应用到源模型上。量化配置编码了关于
如何量化模型的所有信息，例如位数、数据类型、对称性、
scope 名称等。
从本质上讲，量化配置定义为以下类型的命令集合：
_“在 **Scope Z** 下，使用 **ConfigN** 对 **Operator Y** 应用 **量化算法 X**。”_
例如：
_\"使用 **INT8 对称与动态量化**，对 scope **'dense1/'** 下的 **FullyConnected op** 进行**均匀量化**。”_
所有未指定的 op 将保持为 FP32（未量化）。在 TFLite 中，operator 的 scope 被定义为该 op 的输出 tensor 名称，这保留了来自源模型（例如 TF 中的 scope）的分层模型信息。获取 scope 名称的最佳方法是使用 [Model Explorer](https://ai.google.dev/edge/model-explorer) 可视化模型。
目前，量化一个 operator 有三种方法：
* **动态量化（推荐）**：对权重进行量化，而
activation 保持浮点格式，且不经过 AI Edge
Quantizer (AEQ) 处理。Runtime kernel 负责处理这些 activation 的即时量化，这由 `compute_precision=integer` 和
`explicit_dequantize=False` 来标识。
* 优点：减少了模型大小和内存使用。由于
整数计算而提升了延迟。不需要样本数据（校准）。
* 缺点：activation tensor 的即时量化可能会影响模型
质量。并非所有硬件都支持（例如，某些 GPU 和 NPU）。
* **仅权重量化**：仅量化模型权重，而不量化
activation。实际的操作 (op) 计算仍保持浮点格式。
量化后的权重在被输入到 op 之前，会通过
在量化权重和消耗它的 op 之间插入一个 dequantize op 来进行显式反量化。
要启用此功能，需将 `compute_precision` 设置为 `float`，并将
`explicit_dequantize` 设置为 `True`。
* 优点：减少了模型大小和内存使用。不需要样本数据
（校准）。通常具有最佳的模型质量。
* 缺点：由于进行带有
显式反量化的浮点计算，因此没有延迟收益（甚至可能更差）。
* **静态量化**：权重和 activation 均被量化。这
需要一个校准阶段来估算 runtime
tensor (activation) 的量化参数。
* 优点：减少了模型大小、内存使用和延迟。
* 缺点：需要样本数据进行校准。将静态量化
参数（源自校准）施加于 runtime tensor 可能会损害
质量。
通常，我们建议在 CPU/GPU 部署时使用动态量化，在 NPU 部署时使用静态
量化。
我们在 [recipe.py](ai_edge_quantizer/recipe.py) 中包含了常用的配置。这在
[入门 Colab](colabs/getting_started.ipynb) 示例中进行了演示。高级用户可以通过量化器 API 构建自己的配置。
#### 模型验证与准确性基准测试
量化模型不可避免地会引入数值噪声。在调用
`qt.quantize()` 之后，您可以使用内置的 `validate()` 方法验证浮点
基线与量化模型之间的数学失真，该
方法会返回一个单一的 `ComparisonResult` 对象，将节点映射到其误差指标
值。您可以将它们打印出来，或者自动保存为 Model Explorer JSON
文件：
```
# 1. 默认 validation（默认评估 MSE metric）
comparison_results = qt.validate(test_data=sample_data)
print(
"Per-layer metrics:",
comparison_results.get_all_tensor_results(),
)
# 2. 多 metric validation（直接保存所有 metrics 和 validation json 数据）
comparison_results = qt.validate(
test_data=sample_data,
error_metrics=[
quantizer.ValidationErrorMetric.MSE,
quantizer.ValidationErrorMetric.SNR,
],
save_folder='/tmp/'
)
all_results = comparison_results.get_all_tensor_results()
for tensor_name, metrics in all_results.items():
print(
f"Tensor: {tensor_name} "
f"- MSE: {metrics.get(quantizer.ValidationErrorMetric.MSE.value, 0.0):.6f} "
f"- SNR: {metrics.get(quantizer.ValidationErrorMetric.SNR.value, 0.0):.6f}"
)
```
更详细的示例可以在 [quantize_toy_model.py](ai_edge_quantizer/examples/mnist/quantize_toy_model.py) 中找到。
#### 使用 Model Explorer 可视化模型
获取确切的 operator scope 名称，并在视觉上比较基线浮点和量化图之间的 tensor 形状和量化缩放比例的最佳方法是使用 [Model Explorer](https://ai.google.dev/edge/model-explorer)。
要在终端中并排可视化两个导出的 `.tflite` 模型，请运行：
```
model_explorer --models \
"/path/to/baseline_float.tflite,/path/to/quantized_model.tflite"
```
#### 部署
请参阅 [LiteRT 部署文档](https://ai.google.dev/edge/litert/inference) 了解部署量化 LiteRT 模型的方法。
### 高级配置
除了使用 [recipe.py](ai_edge_quantizer/recipe.py) 中的模板外，用户还可以通过多种方式配置和自定义量化配置。例如，用户可以配置配方以实现以下功能：
* 选择性量化（将选定的 op 排除在量化之外）
* 灵活的混合方案量化（混合不同的 precision、计算
precision、scope、op、config 等）
* 4-bit 权重量化
[选择性量化 Colab](colabs/selective_quantization_isnet.ipynb)
展示了其中一些更高级的功能。
有关配置 schema 的详细信息，请参阅 [recipe_manager.py] 中的
`OpQuantizationRecipe`。
对于涉及混合量化的高级用法，以下 API 可能会
有用：
* 使用
[quantizer.py](ai_edge_quantizer/quantizer.py) 中的 `Quantizer:load_quantization_recipe()` 来加载自定义配置。
* 使用
[quantizer.py](ai_edge_quantizer/quantizer.py) 中的 `Quantizer:update_quantization_recipe()` 来扩展或覆盖
配置的特定部分。
### Operator 覆盖范围
下表概述了可用配置允许的设置。
| | | | | | | | | | |
| --- | --- | --- | --- | --- | --- |--- |--- |--- |--- |
| **配置** | | DYNAMIC_WI8_AFP32 | DYNAMIC_WI4_AFP32 | STATIC_WI8_AI16 | STATIC_WI4_AI16 | STATIC_WI8_AI8 | STATIC_WI4_AI8 | WEIGHTONLY_WI8_AFP32 | WEIGHTONLY_WI4_AFP32 |
|activation| num\_bits | None | None | 16 | 16 | 8 | 8 | None | None |
| | symmetric |None | None | TRUE | TRUE | [TRUE, FALSE] | [TRUE, FALSE] | None | None |
| | granularity |None | None | TENSORWISE | TENSORWISE | TENSORWISE | TENSORWISE | None | None |
| | dtype| None | None |INT | INT | INT | INT | None | None |
| weight | num\_bits | 8 | 4 | 8 | 4 | 8 | 4 | 8 | 4 |
| | symmetric | TRUE | TRUE | TRUE | TRUE | TRUE | TRUE | [TRUE, FALSE] | [TRUE, FALSE] |
| | granularity | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] | \[CHANNELWISE, TENSORWISE\] |
| | dtype | INT | INT | INT | INT | INT | INT | INT | INT |
| explicit\_dequantize | | FALSE | FALSE | FALSE | FALSE | FALSE | FALSE | TRUE | TRUE |
| compute\_precision || INTEGER | INTEGER | INTEGER | INTEGER | INTEGER | INTEGER | FLOAT | FLOAT |
**支持量化的 Operator**
| | | | | | | | | |
| --- | --- | --- | --- | --- | --- |--- |--- |--- |
| **配置** | DYNAMIC_WI8_AFP32 | DYNAMIC_WI4_AFP32 | STATIC_WI8_AI16 | STATIC_WI4_AI16 | STATIC_WI8_AI8 | STATIC_WI4_AI8 | WEIGHTONLY_WI8_AFP32 | WEIGHTONLY_WI4_AFP32 |
|FULLY_CONNECTED |
✓
| ✓
| ✓
| ✓
| ✓
| ✓
| ✓
| ✓
|
|CONV_2D | ✓
| | ✓
| ✓
| ✓
| ✓
| ✓
| |
|BATCH_MATMUL | ✓
| | ✓
| | ✓
| | ✓
| |
|EMBEDDING_LOOKUP | ✓
| ✓
| ✓
| | ✓
| ✓
| ✓
| |
|DEPTHWISE_CONV_2D| ✓
| | ✓
| | ✓
| | ✓
| |
|AVERAGE_POOL_2D | | | ✓
| | ✓
| | | |
|RESHAPE | | | ✓
| | ✓
| | | |
|SOFTMAX | | | ✓
| | ✓
| | | |
|TANH | | | ✓
| | ✓
| | | |
|TRANSPOSE | | | ✓
| | ✓
| | | |
|GELU | | | ✓
| | ✓
| | | |
|ADD | | | ✓
| | ✓
| | | |
|CONV_2D_TRANSPOSE| ✓
| | ✓
| | ✓
| | | |
|SUB | | | ✓
| | ✓
| | | |
|MUL | | | ✓
| | ✓
| | | |
|MEAN | | | ✓
| | ✓
| | | |
|RSQRT | | | ✓
| | ✓
| | | |
|CONCATENATION | | | ✓
| | ✓
| | | |
|STRIDED_SLICE | | | ✓
| | ✓
| | | |
|SPLIT | | | ✓
| | ✓
| | | |
|LOGISTIC | | | ✓
| | ✓
| | | |
|SLICE | | | ✓
| | ✓
| | | |
|SELECT | | | ✓
| | ✓
| | | |
|SELECT_V2 | | | ✓
| | ✓
| | | |
|SUM | | | ✓
| | ✓
| | | |
|PAD | | | ✓
| | ✓
| | | |
|PADV2 | | | ✓
| | ✓
| | | |
|MIRROR_PAD | | | ✓
| | ✓
| | | |
|SQUARED_DIFFERENCE | | | | | ✓
| | | |
|MAX_POOL_2D | | | ✓
| | ✓
| | | |
|RESIZE_BILINEAR | | | ✓
| | ✓
| | | |
|RESIZE_NEAREST_NEIGHBOR| | | ✓
| | ✓
| | | |
|GATHER_ND | | | ✓
| | ✓
| | | |
|PACK | | | ✓
| | ✓
| | | |
|UNPACK | | | ✓
| | ✓
| | | |
|DIV | | | ✓
| | ✓
| | | |
|SQRT | | | ✓
| | ✓
| | | |
|GATHER | | | ✓
| | ✓
| | | |
|HARD_SWISH | | | | | ✓
| | | |
|MAXIMUM | | | ✓
| | ✓
| | | |
|REDUCE_MIN | | | ✓
| | ✓
| | | |
|EQUAL | | | ✓
| | ✓
| | | |
|NOT_EQUAL | | | ✓
| | ✓
| | | |
|SPACE_TO_DEPTH | | | | | ✓
| | | |
|RELU | | | ✓
| | ✓
| | | |标签：LiteRT, TensorFlow, 人工智能, 模型优化, 模型量化, 用户模式Hook绕过, 端侧部署, 边缘计算, 逆向工具