📚 课程目录第4章支持的模型格式

4、支持的模型格式：TensorFlow、PyTorch、ONNX、Caffe等主流框架模型的导入方式与注意事项

在RKNN-Toolkit2的实际使用中，模型导入是整个转换链路的第一道关口。很多开发者往往将精力集中在量化调参和精度验证上，却忽略了导入阶段的格式适配与算子映射问题，导致后续步骤反复报错或精度异常。本章将系统梳理TensorFlow、PyTorch、ONNX、Caffe四种主流框架模型的导入方式、内部转换机制以及关键注意事项，帮助你在进入量化环节之前，先构建一个稳健的模型输入管道。

上图展示了RKNN-Toolkit2模型导入的整体架构。从四个主流框架出发，经过解析器完成算子映射与布局转换，生成统一的中间表示（IR），最终导出为板端可执行的.rknn文件。右侧列出了导入阶段最常见的风险点，这些内容将在后续小节中逐一展开。

4.1 核心原理：从框架模型到RKNN中间表示的转换机制

RKNN-Toolkit2的模型导入并非简单的文件格式转换，而是一个涉及计算图解析、算子映射、张量布局调整和权重重排的复杂过程。其底层逻辑可以概括为“前端解析 + 中间表示统一 + 后端优化”的三段式架构。

在前端解析阶段，工具链针对不同框架调用对应的解析器。TensorFlow模型通过GraphDef解析器读取.pb文件中的节点定义；PyTorch模型则依赖TorchScript或ONNX导出后的计算图；ONNX模型直接解析其protobuf格式的图结构；Caffe模型需要同时读取网络结构文件（.prototxt）和权重文件（.caffemodel）。每个解析器都会将原始框架的计算图转换为RKNN自定义的中间表示（IR），该IR以张量流图的形式描述网络拓扑、算子类型、权重数据和数据流关系。

在算子映射阶段，工具链将IR中的每个算子与RKNN内置的算子库进行匹配。RKNN-Toolkit2内置了超过100种常用算子，覆盖卷积、池化、全连接、激活函数、归一化、张量操作等类别。对于无法直接映射的算子，工具链会尝试将其分解为多个基础算子的组合，或者调用自定义算子接口。这一过程直接决定了模型能否成功转换以及转换后的运行效率。

张量布局调整是另一个关键环节。不同框架对张量维度的排列方式存在差异：TensorFlow默认使用NHWC（通道在后）布局，而PyTorch和Caffe使用NCHW（通道在前）布局。RKNN-Toolkit2在导入时会自动检测输入张量的布局，并在IR中标注布局信息，后续在量化或编译阶段根据目标硬件（如RK3588、RK3566等）的NPU特性进行必要的布局转换。

核心要点：模型导入的本质是“计算图语义的等价迁移”。转换过程中，工具链关注的是算子的数学含义、张量的数据流关系以及权重的数值精度，而不是框架的API调用方式。因此，只要模型中的算子都能被RKNN算子库覆盖，无论原始框架是什么，最终生成的.rknn模型在功能上是等价的。

4.2 各框架模型的导入方式与操作流程

RKNN-Toolkit2为每种框架提供了对应的加载接口，统一通过rknn.load_tensorflow()、rknn.load_pytorch()、rknn.load_onnx()、rknn.load_caffe()方法实现。以下分别介绍各框架的具体导入流程和参数配置。

4.2.1 TensorFlow模型导入

TensorFlow模型的导入支持两种格式：冻结的.pb文件（GraphDef格式）和.tflite文件。推荐使用冻结后的.pb文件，因为它包含了完整的计算图和权重数据，不依赖训练环境的变量状态。

from rknn.api import RKNN

rknn = RKNN()
ret = rknn.load_tensorflow(
    tf_pb='model.pb',                # 冻结后的模型文件
    inputs=['input_tensor'],          # 输入节点名称
    outputs=['output_tensor'],        # 输出节点名称
    input_size_list=[[1, 224, 224, 3]]  # 输入形状，注意NHWC布局
)
if ret != 0:
    print('TensorFlow模型加载失败')
    exit(ret)

导入时需要注意：inputs和outputs参数必须与.pb文件中的节点名称完全一致，可以通过graph.get_operations()查看节点列表。输入形状使用NHWC格式，与TensorFlow原生布局保持一致。如果模型包含动态维度（如batch size为None），需要在input_size_list中指定具体数值。

4.2.2 PyTorch模型导入

PyTorch模型的导入相对特殊，RKNN-Toolkit2不直接解析.pt或.pth文件，而是要求先将PyTorch模型导出为ONNX格式，再通过ONNX路径加载。这一设计是因为PyTorch的动态图机制难以直接静态化解析。

# 第一步：在PyTorch环境中导出ONNX
import torch
model = YourModel()
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, 'model.onnx',
    input_names=['input'], output_names=['output'],
    opset_version=11, dynamic_axes=None
)

# 第二步：在RKNN-Toolkit2中加载ONNX
rknn = RKNN()
ret = rknn.load_onnx(model='model.onnx')

导出ONNX时，opset_version建议设置在11-13之间，过高的版本可能导致RKNN解析器不兼容。如果模型包含自定义算子或动态控制流，需要在导出前进行静态化处理，例如使用torch.jit.script或torch.jit.trace。

4.2.3 ONNX模型直接导入

ONNX作为中间格式，是RKNN-Toolkit2支持最完善的导入路径。无论是直接从PyTorch、TensorFlow导出，还是从其他工具链转换得到，只要ONNX模型结构合法，都可以直接加载。

rknn = RKNN()
ret = rknn.load_onnx(model='model.onnx', input_size_list=[[1, 3, 224, 224]])

ONNX导入时，工具链会自动解析模型的输入输出名称和形状。如果模型包含动态轴（如动态batch），需要在input_size_list中指定固定值。建议在导出ONNX时固定所有维度，避免后续量化阶段出现形状推断错误。

4.2.4 Caffe模型导入

Caffe模型的导入需要同时提供网络结构文件（.prototxt）和权重文件（.caffemodel）。RKNN-Toolkit2会解析prototxt中的层定义，并从caffemodel中读取对应的权重数据。

rknn = RKNN()
ret = rknn.load_caffe(
    model='deploy.prototxt',      # 网络结构文件
    proto='model.caffemodel'      # 权重文件
)

Caffe模型导入时，工具链会自动处理其特有的层类型，如Convolution、Pooling、ReLU等。需要注意的是，Caffe的输入布局为NCHW，且deploy.prototxt中必须包含input_shape或input_dim定义，否则需要手动指定输入尺寸。

框架	文件格式	加载接口	输入布局	关键参数
TensorFlow	.pb / .tflite	load_tensorflow()	NHWC	inputs、outputs、input_size_list
PyTorch	.pt → .onnx	load_onnx()	NCHW	opset_version、dynamic_axes
ONNX	.onnx	load_onnx()	NCHW	input_size_list
Caffe	.prototxt + .caffemodel	load_caffe()	NCHW	model、proto

4.3 案例拆解：一个PyTorch分类模型的完整导入过程

下面通过一个具体的图像分类模型案例，展示从PyTorch训练模型到RKNN模型导入的完整闭环，并分析过程中可能遇到的问题。

背景与约束：某AIoT项目需要在RK3588平台上部署一个MobileNetV3-Small图像分类模型，输入为224×224×3的RGB图像，共1000类。开发环境为Ubuntu 20.04，PyTorch 1.13，RKNN-Toolkit2版本为1.6.0。

初始现象：直接使用torch.onnx.export导出ONNX后，在RKNN-Toolkit2中加载时出现“Unsupported operator: aten::hardswish”的错误提示。

分析判断：MobileNetV3使用了HardSwish激活函数，而RKNN-Toolkit2 1.6.0版本的ONNX解析器尚未支持该算子的直接映射。通过查看ONNX模型结构，确认HardSwish节点存在于计算图中。

实际操作步骤：

在PyTorch导出时，将HardSwish替换为等价的数学表达式组合：x * relu6(x + 3) / 6。通过修改模型定义或使用torch.onnx.export的custom_opsets参数，将HardSwish拆解为多个基础算子。
重新导出ONNX，使用opset_version=12，并固定输入形状为[1, 3, 224, 224]。
在RKNN-Toolkit2中加载ONNX模型，设置input_size_list=[[1, 3, 224, 224]]。
执行rknn.build(do_quantization=False)进行预编译，验证模型结构是否完整。
确认无误后，再执行量化构建，生成最终的.rknn文件。

结果对比：替换HardSwish后，模型成功导入并完成量化，在RK3588上的推理精度与PyTorch原始模型相比，Top-1准确率仅下降0.3%（从71.2%降至70.9%），满足项目需求。

风险与代价：手动拆解算子增加了模型修改的工作量，且拆解后的计算图节点数增加，可能导致推理速度略有下降（实测约2%）。但相比更换模型结构或等待工具链更新，这是最稳妥的解决方案。

复盘经验：在模型选型阶段，应提前查阅RKNN-Toolkit2的算子支持列表，避免使用过于新颖或罕见的激活函数。对于常见的HardSwish、Mish等激活函数，建议在导出前进行等价替换。

4.4 常见误区与总结复盘

模型导入阶段的错误往往具有隐蔽性，很多问题在转换时不会立即暴露，而是在板端推理时才显现。以下是实践中总结的高频误区和应对策略。

高风险误区一：忽略输入张量的布局差异。TensorFlow模型使用NHWC布局，而PyTorch/Caffe/ONNX使用NCHW布局。如果混淆了布局，模型虽然能完成转换，但推理结果会完全错误。建议在导入后使用rknn.list_nodes()检查输入张量的形状和布局标注。

高风险误区二：动态形状模型未固定维度。许多PyTorch模型在导出ONNX时保留了动态batch维度，这会导致RKNN在量化阶段无法确定张量形状。务必在导出时使用dynamic_axes=None或指定固定值。

经验提醒：在导入模型后，建议立即执行一次rknn.build(do_quantization=False)进行结构验证。这一步不涉及量化，耗时较短，可以快速发现算子不支持、节点连接错误等问题。确认结构无误后再进行量化，可以大幅减少排障时间。

总结来看，模型导入的成败取决于三个核心要素：算子兼容性（模型中的算子是否被RKNN支持）、布局一致性（输入输出张量的维度排列是否与预期一致）、形状确定性（所有张量维度是否在编译期可推断）。在实际项目中，建议建立标准化的模型导入检查清单：确认框架版本与工具链版本兼容性、检查算子支持列表、固定输入形状、验证输入输出名称、执行预编译验证。通过系统化的流程管理，可以将导入阶段的错误率降低80%以上。