📚 课程目录第4章模型转换流程

4、模型转换流程：PyTorch模型导出ONNX、ONNX模型检查与简化、算能TPU模型转换工具（tpu-mlir）使用、生成INT8/FP16模型

在算能AI盒子（如SE5、SM5系列）上部署YOLO目标检测模型，核心工作并非训练，而是将训练好的PyTorch模型高效、无损地转换为算能TPU可执行的格式。整个转换链路可抽象为“三段式”：PyTorch → ONNX（算法结构标准化）→ ONNX检查与简化（结构净化）→ tpu-mlir（TPU专用编译与量化）。本章将深入拆解这一流程中的关键操作、工具使用细节、量化原理及常见陷阱，帮助工程师建立一条可复现、可排障的模型转换流水线。

4.1 PyTorch模型导出ONNX：核心原理与操作流程

核心原理：PyTorch模型本质是Python类与动态计算图的组合，而TPU编译器无法直接解析Python对象。ONNX（Open Neural Network Exchange）作为一种静态计算图格式，将模型结构、权重、算子类型固化为标准化的protobuf文件。导出过程本质上是“用一组标准算子重新描述前向计算逻辑”，同时将权重参数序列化。对于YOLO系列模型，导出时需特别注意后处理部分（NMS）是否包含在计算图中。通常建议将NMS留在CPU端处理，仅导出主干+检测头（即输出原始预测特征图），这样可保持模型结构简洁，便于后续量化与优化。

操作流程：以YOLOv5为例，导出命令通常如下：

python export.py --weights yolov5s.pt --include onnx --opset 12 --simplify --batch-size 1

关键参数说明：

opset：ONNX算子集版本。tpu-mlir对opset 11~13支持较好，推荐使用12或13。过高版本可能导致部分新算子无法被TPU解析。
dynamic：是否导出动态轴。算能TPU固定batch为1时性能最优，建议导出静态shape（1,3,640,640）。
simplify：导出时自动调用onnx-simplifier进行初步简化。

导出后务必验证输出张量的shape是否符合预期。YOLOv5s在640×640输入下，输出应为三个特征图：(1, 255, 80, 80)、(1, 255, 40, 40)、(1, 255, 20, 20)（255 = 3 anchors × (5 + 80类)）。若shape异常，需检查模型定义中的输出层拼接逻辑。

经验提醒：导出前务必设置model.eval()并关闭梯度。同时，若模型包含自定义算子（如Focus模块），建议在导出前将其重构为标准的Conv+Slice组合，否则ONNX中会出现不兼容节点。

4.2 ONNX模型检查与简化：保障结构纯净度

核心原理：导出的ONNX模型往往包含大量冗余算子（如Identity、Shape、Gather等），这些算子源于PyTorch动态图的自动微分机制。虽然不影响推理正确性，但会增加TPU编译器的解析负担，甚至导致某些算子无法映射到TPU指令。因此，必须进行结构检查与图优化。

检查方法：使用Python脚本加载ONNX模型并验证其完整性：

import onnx
model = onnx.load("yolov5s.onnx")
onnx.checker.check_model(model)  # 结构完整性检查
# 打印所有算子类型
ops = set(node.op_type for node in model.graph.node)
print("算子集合:", ops)

重点关注是否存在以下高风险算子：Loop、If、NonMaxSuppression、RoiAlign。这些算子通常无法被TPU直接支持，需在导出时规避或通过简化工具消除。

简化操作：使用onnx-simplifier进行常量折叠与冗余消除：

python -m onnxsim yolov5s.onnx yolov5s_sim.onnx --overwrite-input-shape 1,3,640,640

简化后的模型应满足：算子数量减少30%以上，且所有节点的输入输出维度均为静态值。可通过Netron可视化工具对比简化前后的计算图结构。

检查项	通过标准	失败处理
算子类型	仅包含Conv、BatchNorm、Relu、Sigmoid、Add、Mul、Reshape、Transpose等基础算子	手动重构模型源码，替换复杂算子
输入输出shape	完全静态，无动态维度	重新导出，固定batch与分辨率
权重精度	FP32或FP16，无INT8混合	统一转为FP32后重新导出
图结构	无孤立节点，无重复子图	使用onnx-simplifier二次清理

风险提示：不要过度简化。某些简化操作会将BatchNorm层折叠进Conv层，这虽然减少了算子数量，但会改变权重分布，对后续INT8量化精度产生微妙影响。建议保留BatchNorm折叠，但务必在量化校准前重新统计激活值范围。

4.3 算能TPU模型转换工具（tpu-mlir）使用详解

核心原理：tpu-mlir是算能科技提供的开源编译器工具链，其作用是将ONNX模型转换为TPU可执行的bmodel二进制文件。转换过程分为两步：model_transform（将ONNX转换为MLIR中间表示）和model_deploy（将MLIR编译为bmodel）。该工具链支持FP32、FP16、INT8三种精度，其中INT8需要提供校准数据集。

环境准备：在Docker容器中运行tpu-mlir是最佳实践。算能官方提供sophgo/tpuc_dev镜像，包含完整的依赖环境。启动命令：

docker run -v $PWD:/workspace -it sophgo/tpuc_dev:latest bash

转换流程：

ONNX转MLIR：

model_transform.py --model_name yolov5s \
  --model_def yolov5s_sim.onnx \
  --input_shapes [[1,3,640,640]] \
  --mean 0.0,0.0,0.0 \
  --scale 0.003921,0.003921,0.003921 \
  --pixel_format rgb \
  --mlir yolov5s.mlir

参数说明：mean和scale用于输入预处理。YOLOv5训练时通常使用0-1归一化，因此scale设为1/255≈0.003921，mean为0。若训练时使用ImageNet标准化（mean=[0.485,0.456,0.406]），则需对应修改。

MLIR转F32 bmodel（用于精度验证）：

model_deploy.py --mlir yolov5s.mlir \
  --quantize F32 \
  --chip bm1684x \
  --output yolov5s_f32.bmodel

MLIR转INT8 bmodel（需校准）：

model_deploy.py --mlir yolov5s.mlir \
  --quantize INT8 \
  --calibration_table yolov5s_cali_table \
  --chip bm1684x \
  --output yolov5s_int8.bmodel

其中calibration_table由run_calibration.py生成，需要准备100~200张代表性图片。

关键参数解析：--chip参数需根据实际AI盒子型号选择。SE5系列对应bm1684，SE7系列对应bm1684x，SM5系列对应bm1688。选错芯片型号将导致编译失败或运行错误。

4.4 生成INT8/FP16模型：量化原理与实战案例

核心原理：FP16模型将权重和激活值从FP32降为FP16，模型体积减半，推理速度提升约1.5倍，精度损失通常小于0.5%。INT8量化则进一步将数据映射到8位整数，推理速度提升3~5倍，但精度损失可能达到1%~3%。YOLO模型对量化敏感度较高，尤其是检测头部分的分类与回归分支，因此需要精细的校准策略。

量化校准流程：

# 准备校准图片列表
echo "cali_images/0001.jpg" > cali_list.txt
echo "cali_images/0002.jpg" >> cali_list.txt
# 运行校准
run_calibration.py --model yolov5s.mlir \
  --dataset cali_list.txt \
  --input_num 100 \
  --calibration_table yolov5s_cali_table

校准集应覆盖不同光照、不同目标尺度、不同背景复杂度的图片。对于目标检测任务，建议从训练集中随机抽取，并确保包含小目标样本。

案例拆解：某智慧园区项目YOLOv5s量化实战

项目背景：需要在SE5盒子（bm1684）上部署YOLOv5s模型，实现园区人员与车辆检测。原始PyTorch模型mAP@0.5为0.892，FP32 bmodel推理帧率仅12FPS，无法满足25FPS的实时要求。

初始现象：直接使用默认参数生成INT8模型后，mAP降至0.847，精度损失4.5%，且小目标（行人）漏检率明显上升。

分析过程：通过对比FP32与INT8模型在各层输出的激活值分布，发现检测头倒数第二层的Conv输出范围达到[-12.3, 11.8]，远超INT8的[-128,127]表示范围，导致严重截断误差。

优化措施：

在导出ONNX前，在检测头前增加一个Clip层，将特征值限制在[-6,6]范围内，减少极端值影响。
调整校准集，增加夜间红外图片与远距离小目标样本，使校准分布更贴近实际场景。
使用--quantize INT8 --fuse_preprocess参数，将图像预处理（缩放、归一化）融合进模型，减少外部干扰。

结果对比：优化后INT8模型mAP回升至0.881，推理帧率提升至38FPS，满足实时性要求。代价是模型体积从14MB增至14.5MB（Clip层引入少量参数），且校准时间增加20分钟。

复盘经验：INT8量化不是简单的“一键转换”，需要结合数据分布、算子特性进行针对性调优。对于YOLO这类anchor-based检测器，优先检查检测头输出的数值范围，必要时进行截断或缩放。

精度模式	模型体积	推理帧率(FPS)	mAP@0.5	适用场景
FP32	28MB	12	0.892	精度优先，无实时要求
FP16	14MB	18	0.889	平衡精度与速度
INT8（默认）	7MB	35	0.847	速度优先，精度可接受
INT8（优化后）	7.5MB	38	0.881	推荐部署方案

常见误区：不要直接使用训练时的验证集作为校准集。验证集通常经过精心筛选，分布过于理想，导致校准后的量化参数在实际场景中失效。应使用独立的、带噪声的真实场景图片。

最后，生成bmodel后务必在目标设备上进行端到端验证。使用bmrt_test工具加载模型并执行推理，检查输出张量的shape与数值范围是否符合预期。若发现输出全零或NaN，优先检查输入预处理参数（mean/scale）是否与训练时一致，其次检查量化表是否包含所有需要校准的层。