3、模型格式入门：ONNX模型导出、TensorFlow模型冻结、PyTorch模型转ONNX踩坑记

各位同学，欢迎来到第三章。

前面我们聊了RKNN工具链的安装和整体流程。今天，咱们来点实际的——模型格式转换。

你想想看，瑞芯微的NPU它不认识PyTorch的.pth，也不认识TensorFlow的.ckpt。它只认自己的RKNN格式。所以，第一步就是把你的模型转成中间格式——ONNX。

ONNX就像模型界的“普通话”。不管你是用PyTorch、TensorFlow还是其他框架，先转成ONNX，再转RKNN，这是最稳妥的路径。我在项目中试过直接转，嗯，踩坑概率会高不少。

3.1 ONNX模型导出：从PyTorch到ONNX

先说PyTorch。PyTorch转ONNX其实有官方API，叫torch.onnx.export。但这里面的门道不少。

我直接上代码吧。这是我最常用的导出脚本：

import torch
import torchvision.models as models

# 1. 加载模型，设为eval模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造一个样例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 样例输入
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 导出参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)
print("ONNX导出完成")

这里有个细节——opset_version。我建议用11或12。为什么？因为RKNN工具链对这两个版本的支持最成熟。我曾经试过opset=15，结果RKNN转换时报错说某个算子不支持，折腾了半天。

3.2 TensorFlow模型冻结：从SavedModel到PB

TensorFlow这边，情况稍微复杂一点。因为TensorFlow有多个版本，模型保存格式也变了好几次。

如果你用的是TensorFlow 1.x，模型通常保存为checkpoint。你需要先冻结成.pb文件。冻结的意思是把计算图和权重打包成一个文件。

我给大家一个通用的冻结脚本：

import tensorflow as tf
from tensorflow.python.framework import graph_io

# 1. 加载meta图和checkpoint
sess = tf.Session()
saver = tf.train.import_meta_graph('model.meta')
saver.restore(sess, 'model')

# 2. 获取输入输出节点名称
graph = tf.get_default_graph()
input_tensor = graph.get_tensor_by_name('input:0')
output_tensor = graph.get_tensor_by_name('output:0')

# 3. 冻结
frozen_graph = tf.graph_util.convert_variables_to_constants(
    sess,
    sess.graph_def,
    ['output']  # 输出节点名称
)

# 4. 保存
graph_io.write_graph(frozen_graph, '.', 'frozen_model.pb', as_text=False)
print("冻结完成")

如果你用的是TensorFlow 2.x，那就简单多了。直接用tf.saved_model.save保存，然后转ONNX就行。

3.3 PyTorch转ONNX踩坑记

这部分我得多说几句。PyTorch转ONNX，看着简单，实际坑不少。我把常见的坑列出来，你们遇到类似问题可以直接翻到这里。

坑一：动态图与静态图的不匹配

PyTorch是动态图，支持if-else、for循环这些控制流。但ONNX是静态图，不支持动态控制流。

举个例子，如果你的模型里有这样的代码：

if x.sum() > 0:
    y = self.branch_a(x)
else:
    y = self.branch_b(x)

导出时，ONNX只会记录其中一个分支。另一个分支就丢了。怎么办？

我的做法是：用torch.where或者torch.max这类静态操作替代if-else。实在不行，就改模型结构。

坑二：算子不支持

有些PyTorch算子，ONNX没有对应的实现。比如torch.einsum、torch.topk在某些版本下就不支持。

我记得有一次，我导出一个注意力机制的模型，里面用了torch.einsum。导出时报错：

RuntimeError: Exporting the operator 'einsum' to ONNX opset version 11 is not supported.

解决方案？要么换算子，要么升级opset版本。我最后换成了torch.matmul和torch.transpose的组合。

坑三：动态轴设置错误

如果你需要支持动态batch size，就得设置dynamic_axes。但很多人设置错了。

比如，你的模型输入是[batch, 3, 224, 224]，你只设置了batch维度是动态的。但输出可能也有batch维度，你忘了设置。结果就是，推理时batch=1没问题，batch=4就报错。

坑四：模型没有设为eval模式

这个坑我踩过不止一次。如果你在导出时忘了加model.eval()，模型里的Dropout和BatchNorm会处于训练模式。导出的ONNX权重就不对。

结果就是，ONNX推理结果和PyTorch推理结果对不上。我当时排查了半天，最后发现是这个问题。

3.4 验证ONNX模型

导出ONNX后，一定要验证。我一般用这个脚本：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 检查模型结构
onnx_model = onnx.load('resnet18.onnx')
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过")

# 2. 用ONNX Runtime推理
session = ort.InferenceSession('resnet18.onnx')
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 构造输入
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
result = session.run([output_name], {input_name: dummy_input})
print("ONNX推理完成，输出shape:", result[0].shape)

这一步能验证两件事：一是模型结构没问题，二是推理能跑通。如果这一步都过不了，就别想着转RKNN了。

3.5 总结

好了，这一章的内容就这些。总结一下：

• PyTorch转ONNX：用torch.onnx.export，注意opset版本和eval模式
• TensorFlow转ONNX：1.x先冻结成.pb，2.x用SavedModel，然后统一转ONNX
• 踩坑重点：动态图静态图不匹配、算子不支持、动态轴设置、eval模式
• 验证：用onnxruntime跑一遍，确认结果正确

下一章，我们终于要进入正题了——用RKNN工具链把ONNX转成RKNN格式。到时候我会详细讲量化、优化和常见错误。咱们下章见。