第4章模型转换工具链

第4章：模型转换工具链：ONNX导出、RKNN工具链、昇腾CANN、地平线OE

模型训练好了，就像菜做好了。但要让NPU这个“胃”能消化，还得经过一道“烹饪”工序——模型转换。说白了，就是把PyTorch、TensorFlow这些框架的模型，变成NPU能读懂的文件。

我这些年跟各种NPU打交道，发现一个规律：模型转换花的时间，往往比模型训练还多。为什么？因为每个NPU都有自己的“口味偏好”。今天我们就来聊聊主流的几个工具链。

4.1 ONNX：模型转换的“通用语言”

ONNX（Open Neural Network Exchange）是微软和Meta搞出来的中间格式。我个人习惯，所有模型先转ONNX，再转目标平台格式。为什么？因为ONNX的生态最成熟，调试工具也多。

4.1.1 导出ONNX的基本流程

以PyTorch为例，导出ONNX其实就几行代码：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = torch.load('model.pth')
model.eval()

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    },
    opset_version=11
)

这里有个坑，我踩过好几次。opset_version不是越高越好。有些NPU工具链只支持到opset 11或12，你用了15，转换时就报错。我建议先查目标平台的文档，再决定版本。

我的经验：导出ONNX后，先用onnxruntime跑一遍推理，确认精度没问题。这一步能过滤掉80%的转换问题。

4.1.2 ONNX的常见问题

导出ONNX时，最头疼的是动态形状。比如NLP模型，输入长度不固定。你想想看，NPU最讨厌的就是动态形状，它喜欢固定尺寸。

我曾经有个项目，模型输入是变长的语音特征。导出ONNX时，我用了dynamic_axes，结果在RKNN上死活跑不起来。后来改成固定长度padding，问题就解决了。

4.2 RKNN工具链：瑞芯微的“看家本领”

RKNN是瑞芯微NPU的模型格式。说实话，RKNN工具链是我用过最“亲民”的。它把转换、量化、推理都集成在一个包里，上手很快。

4.2.1 RKNN转换流程

先装工具链：

pip install rknn-toolkit2

然后写转换脚本：

from rknn.api import RKNN

# 创建RKNN对象
rknn = RKNN()

# 配置量化参数
rknn.config(
    mean_values=[[123.675, 116.28, 103.53]],
    std_values=[[58.395, 57.12, 57.375]],
    target_platform='rk3588',
    quantize=True
)

# 加载ONNX模型
ret = rknn.load_onnx(model='model.onnx')
if ret != 0:
    print('加载ONNX失败')
    exit(-1)

# 构建RKNN模型
ret = rknn.build(do_quantization=True, dataset='dataset.txt')
if ret != 0:
    print('构建RKNN失败')
    exit(-1)

# 导出RKNN文件
rknn.export_rknn('model.rknn')

# 释放资源
rknn.release()

关键点：dataset.txt文件里放的是量化用的校准图片路径。我建议至少放100张，覆盖各种场景。太少的话，量化后精度会掉得很厉害。

4.2.2 量化避坑指南

我曾经有个模型，量化后精度从98%掉到85%。排查了半天，发现是激活值分布太宽。RKNN默认用int8量化，如果激活值范围在[-10, 10]，量化精度就很高。但如果范围在[-100, 100]，量化误差就大了。

解决办法？在模型里加个BatchNorm层，或者用per-channel量化。RKNN支持per-channel，但需要手动配置。

4.3 昇腾CANN：华为的“重型武器”

昇腾的CANN（Compute Architecture for Neural Networks）工具链，说实话，学习曲线比较陡。但一旦上手，功能确实强大。

4.3.1 模型转换工具：ATC

ATC（Ascend Tensor Compiler）是CANN的核心转换工具。用法如下：

atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --input_shape="input:1,3,224,224" \
    --soc_version=Ascend310P3 \
    --precision_mode=force_fp16

这里要注意soc_version参数。昇腾有310、310P、910等不同芯片，参数不能混用。我刚开始用的时候，总搞不清该填什么，后来发现用npu-smi info命令就能查到。

4.3.2 精度调优

CANN支持混合精度推理。你可以指定某些层用fp16，某些层用int8。我个人习惯，先全量fp16跑一遍，看精度损失。如果损失在1%以内，就保持fp16。如果损失大，再逐层排查。

注意：CANN的算子支持列表是有限的。有些自定义算子，比如自己写的C++算子，CANN可能不支持。我建议在模型设计阶段，就避开这些算子。

4.4 地平线OE：国产新秀

地平线的OE（Open Explorer）工具链，我接触得晚一些。但用下来感觉，它在模型优化方面做得不错。

4.4.1 OE的转换流程

OE的转换脚本长这样：

from horizon_tc_ui import HB_ONNXConverter

# 创建转换器
converter = HB_ONNXConverter()

# 设置模型路径
converter.model_parameters.onnx_model = 'model.onnx'

# 设置输入信息
converter.model_parameters.input_layout = 'NHWC'
converter.model_parameters.input_shape = '1x224x224x3'

# 设置校准数据
converter.model_parameters.cal_data_path = './calibration_data'

# 开始转换
converter.convert()

# 导出bin文件
converter.export('model.bin')

地平线有个特点，它强制要求输入是NHWC格式。而PyTorch默认是NCHW。所以转换前，你得先调整数据格式。我一般会在ONNX导出时，用permute算子把NCHW转成NHWC。

4.4.2 模型优化技巧

OE工具链自带一个模型分析器，能告诉你哪些层是瓶颈。我记得有个项目，模型里有个大卷积层，耗时占了40%。OE建议我用分组卷积替换，结果推理速度提升了30%。

嗯，这里要注意，OE的优化建议不一定都适用。你得结合自己的场景判断。比如它建议用int8量化，但你的模型对精度要求极高，那就得权衡一下。

4.5 工具链对比总结

工具链	适用平台	量化支持	学习难度	我的评价
ONNX	通用	有限	低	必备中间格式
RKNN	瑞芯微	int8/16	低	上手最快
CANN	昇腾	fp16/int8	高	功能最全
OE	地平线	int8	中	优化能力强

选哪个工具链？说白了，取决于你的硬件平台。如果你用瑞芯微的芯片，那就用RKNN。如果用昇腾，那就学CANN。没有最好的，只有最合适的。

最后说一句，模型转换是个细致活。我建议你每次只改一个参数，然后对比精度和速度。这样出了问题，能快速定位。别像我刚开始那样，一次改一堆参数，结果出问题了都不知道是哪步引起的。

📚 课程目录