第 6 章 TorchAOONNX RuntimeTFLite

第 6 章量化通路：PyTorch / ONNX / TFLite

前五章我们从第一性原理手写了量化器。但真实项目不会自己造轮子——你要在 PyTorch、ONNX Runtime、TensorFlow Lite 这些成熟工具链里选一条路，把量化模型送进推理引擎。这一章用同一批模型（ResNet-18、BERT、MobileNetV2、TinyLlama）走完三条主流通路，并回答一个容易被忽略的问题：量化后的模型，在两条不同通路上跑出来的结果一不一样？

原书仓库 ch6/ 目录的脚本：ch6_pytorch_torchao_path.py、ch6_onnx_export_path.py、ch6_tf_mot_path.py、ch6_verify_equivalence.py。

6.1 PyTorch 通路：TorchAO

ch6_pytorch_torchao_path.py 演示了 PyTorch 官方量化库 TorchAO（PyTorch Architecture Optimization）的用法。它对 ResNet、BERT、TinyLlama 三种模型跑了一套量化实验：

# 来自 ch6/ch6_pytorch_torchao_path.py（结构）
def run_torchao_experiment(model, config_name, ...):
    # config_name 例如: "int8_weight_only" / "int4_weight_only"
    # 1. 用 torchao.quantization.quantize_ 对模型做量化
    # 2. 评估精度、模型大小、延迟
    ...

def inspect_torchao_weights(model, config_name):
    # 检查量化后的权重是什么形态（观察 scale 等元数据）
    ...

def verify_save_load_roundtrip(model, ...):
    # 验证量化模型能保存、能加载、数值不变
    ...

TorchAO 的核心用法一句话：quantize_(model, config)，用配置对象指定量化方案（int8_weight_only、int4_weight_only 等），一步到位。它帮你处理了：

权重量化（weight-only）与激活量化（dynamic quantization）的区别；
保存/加载的 roundtrip（量化后的模型要用配套方式序列化，否则反序列化后精度/数值对不上）；
对 TinyLlama 这种 LLM 也能直接套用（这是第 7 章 LLM 量化在生产侧的常见入口之一）。

6.2 ONNX 通路：动态 vs 静态量化

ch6_onnx_export_path.py 走的是 PyTorch → ONNX → ONNX Runtime 的路。它用 ResNet-18 做了完整对比：

# 来自 ch6/ch6_onnx_export_path.py（结构）
def export_resnet_onnx(save_dir):
    # PyTorch 模型导出为 ONNX（FP32）
    ...

def quantize_resnet_dynamic(fp32_path, save_dir):
    # ONNX Runtime 动态量化：权重量化，激活按运行时数据动态计算 scale
    ...

def quantize_resnet_static(fp32_path, save_dir):
    # ONNX Runtime 静态量化：权重 + 激活都量化，激活 scale 来自校准集
    ...

def run_mixed_precision_analysis(fp32_path, save_dir):
    # 混合精度分析：哪些算子量化收益大、哪些敏感
    ...

def inspect_quantizable_nodes(fp32_path):
    # 检查 ONNX 图里哪些节点可以量化
    ...

三个关键概念：

动态量化（dynamic）：只量化权重，激活在运行时"看到"真实输入再算 scale。简单、不需要校准数据，但激活的量化/反量化在运行时多一次计算，加速有限。
静态量化（static）：权重和激活都量化，激活的 scale 用校准集提前定死。推理时没有动态开销，加速最明显，但需要第 4 章的校准流程。
混合精度：不是所有算子都值得量化（例如某些 LayerNorm、Softmax 保持 FP32），run_mixed_precision_analysis 帮你找出收益和风险的平衡点。

ONNX 的量化图会在算子间插入 QDQ（Quantize-DeQuantize）节点——这是 ONNX Runtime / TensorRT 推理引擎识别量化算子的标准方式，第 9 章会再见到它。

6.3 TFLite 通路：MobileNetV2 + BERT

ch6_tf_mot_path.py 用 TensorFlow Lite（TFLite）把 MobileNetV2 和 BERT 走了一遍四种量化模式：

# 来自 ch6/ch6_tf_mot_path.py（结构）
def convert_fp32(saved_model_path, output_path):
    # 基线：不量化
    ...

def convert_dynamic_range(saved_model_path, output_path):
    # 动态范围量化（权重量化 + 激活运行时 scale）
    ...

def convert_full_integer(saved_model_path, output_path, representative_dataset):
    # 全整数量化：权重 + 激活都量化，需要代表性数据集
    ...

def convert_float16(saved_model_path, output_path):
    # FP16 量化：权重转 FP16（半精度，不是整数量化）
    ...

def make_mobilenet_representative_dataset(images, num_calibration=200):
    # 代表性数据集生成器：喂给全整数量化做校准
    ...

TFLite 的 representative_dataset 就是第 4 章"校准集"在 TFLite 世界里的名字——它是一段 Python 生成器，converter 会在转换过程中调用它来统计激活范围。四种模式的取舍：

模式	权重	激活	需要校准数据	典型场景
FP16	FP16	FP32	否	有 FP16 加速的移动 GPU
Dynamic range	INT8	运行时 FP32→INT8	否	快速减体积
Full integer	INT8	INT8	是	CPU/Edge TPU 全整型加速
Float32	FP32	FP32	否	基线

全整数量化是移动端/边缘端的主流选择（只有全整型才能跑 Edge TPU / 各种 NPU 加速器），代价是必须准备代表性数据集。

6.4 跨通路验证：数值等价性

ch6_verify_equivalence.py 是这一章最被低估的脚本。它提出了一个尖锐的问题：同一个 FP32 模型，用 TorchAO 量化和用 ONNX 量化，结果一样吗？ 答案是：不一定。不同工具链的 scale 计算方式、QDQ 图结构、算子融合策略都不同，量化误差会以不同方式传播。

# 来自 ch6/ch6_verify_equivalence.py（核心）
def compute_equivalence_metrics(baseline, variant, labels=None):
    # 计算基准输出与量化输出的等价性指标
    # 例如：最大绝对误差、平均绝对误差、输出张量相关性
    ...

def torchao_resnet_int8wo(images_np, model_fp32):
    # TorchAO 量化后的 ResNet 推理
    ...

def ort_inference(onnx_path, feed_dict):
    # ONNX Runtime 推理
    ...

# 用同一批输入图片/文本，分别跑：
#   FP32 baseline、TorchAO INT8、ONNX INT8
# 比较三者的输出差异

为什么要做这件事：

可复现性：如果你的论文/报告说"INT8 精度掉了 0.5%"，必须说清是哪个工具链量化的——换个工具链可能掉 1%。
迁移风险：从一个推理引擎换到另一个（例如 TorchAO 原型 → ONNX 生产），量化行为变了，必须重新验证。
排查手段：当线上精度异常时，先确认"是不是量化通路切换导致的"，再深入模型本身。

验证方法是朴素的：同一批输入，分别跑 FP32 baseline 和各量化版本，逐样本比较输出张量的差异（最大/平均绝对误差、相关性），必要时再比较下游任务指标（分类准确率、困惑度）。

6.5 本章小结

三条通路的选型逻辑：

graph LR
    A[模型] --> B[PyTorch TorchAO]
    A --> C[ONNX Runtime]
    A --> D[TFLite]
    B --> E[原型/实验/PyTorch 生态内推理]
    C --> F[生产部署 CPU/GPU/TRT/OpenVINO]
    D --> G[移动端/嵌入式/Edge TPU]
    E --> H[跨通路数值等价性验证]
    F --> H
    G --> H

TorchAO：最贴近 PyTorch，适合实验和 PyTorch 生态内推理，quantize_ 一步到位。
ONNX Runtime：动态（快、免校准）vs 静态（更快、需校准）vs 混合精度，QDQ 图是推理引擎的标准语言。
TFLite：FP16 / 动态 / 全整型三档，全整型是移动端和 NPU 的唯一选择，需代表性数据集。
跨通路验证：同一模型在不同工具链下量化行为不同，上线前必须做数值等价性验证。

下一章是全书的重头戏——LLM 量化。第 7 章会看到为什么大模型量化和小模型量化是两回事，以及 LLM.int8()、GPTQ、AWQ 这些名字背后的真正算法。

作者: itech001 来源: 公众号：AI人工智能时代(the-ai-era) 网站: https://www.theaiera.top/ 关注每日最新AI新闻和技术博客，主页有更多的文章的AI 技术参考：https://www.theaiera.top

本文首发于 AI人工智能时代，转载请注明出处。

第 6 章 量化通路：PyTorch / ONNX / TFLite

第 6 章 量化通路：PyTorch / ONNX / TFLite

6.1 PyTorch 通路：TorchAO

6.2 ONNX 通路：动态 vs 静态量化

6.3 TFLite 通路：MobileNetV2 + BERT

6.4 跨通路验证：数值等价性

6.5 本章小结

📑 高效 AI 实战：量化与快速推理

第 6 章量化通路：PyTorch / ONNX / TFLite

第 6 章量化通路：PyTorch / ONNX / TFLite