AI模型效率工具包（AIMET）

AIMET 是一个用于量化训练好的机器学习模型的软件工具包。

AIMET通过减少计算负载和内存占用，提升深度学习模型的运行时性能。使用AIMET量化的模型能够降低内存占用，从而更易于部署在手机或笔记本电脑等边缘设备上。

AIMET采用后训练量化和微调技术，在量化和压缩过程中最大限度地减少精度损失。AIMET支持ONNX和PyTorch框架中的模型。

AIMET专为与PyTorch和ONNX模型配合使用而设计。

您可以在Qualcomm AI Hub Models中找到使用AIMET量化的模型——这是一个优化和量化的模型集合。

为什么选择AIMET？

• 先进的量化技术：使用整数运行时进行推理的速度显著快于浮点运行时。例如，模型在高通Hexagon DSP上的运行速度比在高通Kyro CPU上快5到15倍。此外，8位精度模型的内存占用仅为32位精度模型的四分之一。然而，在量化机器学习模型时，保持模型精度往往颇具挑战性。AIMET通过诸如无数据量化等创新技术解决了这一问题，在多个流行模型上实现了最先进的INT8结果。
• 支持先进的模型压缩技术，使模型在推理时运行更快、所需内存更少。
• AIMET旨在自动化神经网络的优化，避免耗时且繁琐的手动调整。AIMET还提供了用户友好的API，允许用户直接从其PyTorch流水线中调用。

更多详情请访问AIMET的Github Pages。

快速入门

aimet-onnx和aimet-torch已在PyPI上发布。

请参阅我们的快速入门指南，以开始使用最新的AIMET软件包。

从源代码构建

要从源代码构建最新的AIMET代码，请参阅在Docker环境中从源代码构建、安装并运行AIMET

支持的功能

后训练量化(PTQ)

查看指南，开始使用PTQ技术。

下表总结了从基础技术如“校准”到高级技术如“SeqMSE”和“自适应舍入(AdaRound)”等，您可在AIMET中使用的技术。

技术	ONNX	PyTorch	作用
校准	✅	✅	计算量化参数
AdaRound	✅	✅	对量化后的权重进行舍入
SeqMSE	✅	✅	优化每一层的编码
BatchNorm折叠	✅	✅	将BatchNorm折叠，以弥合仿真与目标设备之间的差距
跨层均衡	✅	✅	重新缩放权重以减少范围失衡
BatchNorm重新估计	✅	✅	重新估计BatchNorm统计信息
AdaScale	✅	✅	优化量化后的权重
OmniQuant	❌	✅	优化量化后的权重
SpinQuant	❌	✅	优化量化后的权重

量化感知训练(QAT)

AIMET通过aimet-torch支持量化感知训练(QAT)。

如果您希望同时使用QAT以及AIMET中的一些高级PTQ技术，我们建议采用以下工作流程：

详细QAT指南请见此处。

模型压缩

• 空间SVD：一种张量分解技术，可将大型层拆分为两个较小的层。
• 通道剪枝：移除层中冗余的输入通道，并重建层权重。
• 逐层压缩比选择：自动选择模型中每一层的压缩程度。

可视化

• 权重范围：直观检查模型是否适合应用跨层均衡技术，以及应用该技术后的效果。
• 逐层压缩敏感度：直观了解模型中任意一层对压缩的敏感程度。

结果

AIMET 可以将现有的 32 位浮点模型量化为 8 位定点模型，而几乎不会损失精度，也无需对模型进行微调。

DFQ

将 DFQ 方法应用于 MobileNet-v2 和 ResNet-50 等几种流行网络时，即使量化到 8 位，准确率的损失也始终低于 0.9%，且整个过程完全自动化，无需任何训练数据。

模型	FP32	INT8 模拟
MobileNet v2 (top1)	71.72%	71.08%
ResNet 50 (top1)	76.05%	75.45%
DeepLab v3 (mIOU)	72.65%	71.91%

AdaRound（自适应舍入）

ADAS 目标检测

对于这个难以量化至 8 位精度的 ADAS 目标检测模型，AdaRound 能够将准确率恢复到与 FP32 准确率相差不到 1% 的水平。

配置															mAP - 平均精度均值
FP32															82.20%
最接近舍入（INT8 权重，INT8 激活值）															49.85%
AdaRound（INT8 权重，INT8 激活值）															81.21%

DeepLabv3 语义分割

对于 DeepLabv3 语义分割模型等部分模型，AdaRound 甚至可以将模型权重量化至 4 位精度，而准确率几乎没有明显下降。

配置															mIOU - 平均交并比
FP32															72.94%
最接近舍入（INT4 权重，INT8 激活值）															6.09%
AdaRound（INT4 权重，INT8 激活值）															70.86%

循环神经网络模型的量化

AIMET 支持对循环神经网络模型（RNN、LSTM、GRU）进行量化模拟和量化感知训练（QAT）。利用 AIMET 中的 QAT 功能，带有双向 LSTM 的 DeepSpeech2 模型可以被量化至 8 位精度，同时保持极小的准确率损失。

DeepSpeech2 （使用双向 LSTM）	词错误率
FP32	9.92%
INT8	10.22%

模型压缩

AIMET 还能够显著压缩模型。对于 Resnet-50 和 Resnet-18 等流行模型，通过空间 SVD 加上通道剪枝进行压缩，可以在保留原始未压缩模型约 1% 准确率的情况下，实现 50% 的 MAC（乘累加）减少。

模型	未压缩模型	50% 压缩后的模型
ResNet18 (top1)	69.76%	68.56%
ResNet 50 (top1)	76.05%	75.75%

资源

• 文档主页
• API 参考
• 讨论论坛
• Slack
• 教程视频
• 示例代码

贡献

感谢您对参与 AIMET 贡献的兴趣！请阅读我们的 贡献页面，了解更多关于贡献新功能或修复 bug 的信息。我们期待您的参与！

团队

AIMET 致力于成为一个由高通创新中心公司维护的社区驱动项目。

许可证

AIMET 采用 BSD 3 条款“新”或“修订版”许可证。更多详情请参阅 LICENSE。

AIMET 快速上手指南

AIMET (AI Model Efficiency Toolkit) 是由高通创新中心开发的软件工具包，旨在通过量化（Quantization）和压缩技术优化训练好的深度学习模型。它能显著减少模型的计算负载和内存占用，同时利用先进的后训练量化（PTQ）和量化感知训练（QAT）技术最小化精度损失，非常适合将模型部署到手机、笔记本等边缘设备。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04/22.04)。Windows 和 macOS 支持有限，建议通过 Docker 运行。
• Python 版本: Python 3.8 - 3.10。
• 深度学习框架:
  • PyTorch: 支持主流版本（如 1.13, 2.0, 2.1 等，需与 AIMET 版本对应）。
  • ONNX: 需安装 onnx 库。
• 编译器: GCC/G++ (若需从源码编译)。
• Docker (可选但推荐): 为了规避依赖冲突，官方强烈建议使用提供的 Docker 镜像进行构建和运行。

注意：AIMET 主要针对 Qualcomm 硬件（如 Hexagon DSP）进行了深度优化，但在通用 CPU/GPU 上也可进行量化仿真和验证。

安装步骤

AIMET 提供 PyPI 预编译包，也支持从源码构建。对于大多数用户，推荐使用 pip 安装。

方法一：通过 PyPI 安装（推荐）

根据您的模型框架选择安装 aimet-torch 或 aimet-onnx。

安装 PyTorch 版本：

pip install aimet-torch

安装 ONNX 版本：

pip install aimet-onnx

国内加速提示：如果下载速度较慢，可使用清华或阿里镜像源：

pip install aimet-torch -i https://pypi.tuna.tsinghua.edu.cn/simple

方法二：从源码构建（高级用户）

如果您需要最新的功能或特定的定制版本，建议在 Docker 环境中从源码构建。

1. 克隆仓库：

   git clone https://github.com/quic/aimet.git
   cd aimet
   

2. 按照官方文档在 Docker 中执行构建脚本（具体命令请参考项目根目录下的 build.sh 或官方文档 "Build from source" 章节）：

   # 示例：进入构建环境并执行构建（具体参数需参考最新文档）
   docker run --rm -it -v $(pwd):/workspace aimet-build-env bash
   ./build.sh
   

基本使用

以下是一个使用 AIMET 对 PyTorch 模型进行后训练量化 (PTQ) 的最简示例。该流程展示了如何加载模型、应用量化模拟以及导出量化模型。

示例：PyTorch 模型 INT8 量化

此示例演示如何使用 AutoQuantV2 API 自动选择最佳量化方案。

import torch
from aimet_common.defs import QuantScheme
from aimet_torch.auto_quant_v2 import AutoQuantV2
from aimet_torch.model_preparer import prepare_model

# 1. 加载预训练的 PyTorch 模型
# 假设您有一个标准的 torchvision 模型
from torchvision.models import resnet18
model = resnet18(pretrained=True)
model.eval()

# 2. 准备模型 (处理动态控制流等复杂结构)
# 注意：某些简单模型可能不需要此步骤，但推荐作为标准流程
input_shape = (1, 3, 224, 224)
dummy_input = torch.randn(input_shape)
model_prepared = prepare_model(model)

# 3. 定义评估函数和校准数据加载器
# eval_function: 接收量化后的模型并返回精度分数
def eval_function(model, early_stopping=False):
    # 在此处插入您的验证逻辑
    # 为演示目的，直接返回一个占位分数
    return 0.75 

# data_loader: 用于校准的少量无标签数据 (通常 100-1000 张图像)
# 此处仅为伪代码，实际需替换为真实的 DataLoader
# from torch.utils.data import DataLoader
# calib_data_loader = DataLoader(...) 

# 4. 配置并运行 AutoQuantV2
# AutoQuantV2 会自动尝试不同的量化技术（如 AdaRound, CLE 等）以找到精度损失最小的方案
auto_quant = AutoQuantV2(
    model_prepared,
    dummy_input=dummy_input,
    eval_callback=eval_function,
    # calib_data_loader=calib_data_loader, # 取消注释并传入真实数据
    allowed_accuracy_drop=0.01 # 允许的最大精度损失
)

# 执行量化
try:
    quantized_model, encoding_stats, results = auto_quant.run()
    print("量化完成！")
    print(f"最终精度: {results['accuracy']}")
except Exception as e:
    print(f"量化过程中发生错误：{e}")

# 5. (可选) 导出量化模型和编码文件
# 导出的模型可用于后续在 Qualcomm 目标设备上运行
# from aimet_torch.quantsim import QuantizationSimModel
# sim = QuantizationSimModel(quantized_model, dummy_input)
# sim.export(path='./output', filename_prefix='resnet18_quantized', dummy_input=dummy_input)

核心功能简述

• Post-Training Quantization (PTQ): 无需重新训练，仅需少量校准数据即可将 FP32 模型转换为 INT8，支持 AdaRound, SeqMSE, Cross Layer Equalization 等高级技术。
• Quantization Aware Training (QAT): 在微调过程中模拟量化噪声，进一步恢复精度。
• 模型压缩: 支持空间 SVD 分解和通道剪枝，可大幅减少模型参数量和计算量 (MACs)。

更多详细 API 用法和进阶教程，请访问 AIMET 官方文档。