2. TensorFlow模型转换与量化

任务目标

本案例将完成一个TensorFlow 2.x训练好的MNIST手写数字分类模型（.h5格式）到RKNN格式的完整转换流程，并应用RKNN-Toolkit2的默认量化配置（int8），最终生成可在瑞芯微NPU上部署的mnist.rknn文件。

环境依赖

  • RKNN-Toolkit2（版本≥1.6.0），安装方式参考官方文档
  • TensorFlow 2.x（仅用于生成/验证.h5模型，转换过程本身不依赖TF）
  • Python 3.8+，numpy
  • 测试平台：PC（x86_64）或RK3588等ARM开发板

实现步骤

整个转换流程分为三步：准备TensorFlow模型、编写转换脚本、执行转换并验证输出。默认量化模式下，RKNN-Toolkit2会自动对权重和激活值进行int8量化，无需手动指定量化方案。

关键完整代码

以下脚本完成从.h5到.rknn的完整转换，并包含一个最小验证逻辑（推理随机输入检查输出shape）：

import numpy as np
from rknn.api import RKNN

# 1. 创建RKNN对象（verbose=True可打印详细日志）
rknn = RKNN(verbose=True)

# 2. 配置量化参数（默认配置即int8量化）
#    mean_values/std_values 需与训练时数据预处理一致（MNIST: 像素值0~255，归一化到0~1）
rknn.config(
    mean_values=[[0]],      # 单通道灰度图
    std_values=[[255]],     # 除以255实现归一化
    target_platform='rk3588'  # 根据实际芯片选择，可选 rk3566/rk3588/rk3399pro 等
)

# 3. 加载TensorFlow模型
#    input_size_list 需与模型输入shape一致（NHWC格式）
ret = rknn.load_tensorflow(
    tf_pb='mnist.h5',          # 实际加载时需为.pb或.tflite，见下方说明
    inputs=['input_1'],        # 需与模型输入节点名一致
    outputs=['dense_1'],       # 需与模型输出节点名一致
    input_size_list=[[1, 28, 28, 1]]
)
if ret != 0:
    print('模型加载失败')
    exit(ret)

# 4. 构建RKNN模型（执行默认int8量化）
ret = rknn.build(do_quantization=True, dataset='dataset.txt')
if ret != 0:
    print('构建失败')
    exit(ret)

# 5. 导出RKNN模型
ret = rknn.export_rknn('mnist.rknn')
if ret != 0:
    print('导出失败')
    exit(ret)

# 6. 最小验证：加载模型并推理随机输入
rknn.release()
rknn = RKNN()
rknn.load_rknn('mnist.rknn')
rknn.init_runtime()

# 模拟一张28x28的随机输入（实际使用时需替换为真实图像）
test_input = np.random.randint(0, 255, (1, 28, 28, 1), dtype=np.uint8)
outputs = rknn.inference(inputs=[test_input])
print('推理输出shape:', outputs[0].shape)  # 预期为 (1, 10)
print('输出类型:', outputs[0].dtype)        # 预期为 int8 或 float32（取决于量化配置）

rknn.release()

关键说明与适配

  • 模型格式转换：RKNN-Toolkit2的load_tensorflow接口仅支持.pb（冻结图）或.tflite格式。若你只有.h5文件，需先执行以下转换（在TensorFlow环境中运行）：
import tensorflow as tf

# 加载.h5模型并转换为冻结图.pb
model = tf.keras.models.load_model('mnist.h5')
tf.saved_model.save(model, 'saved_model')
!python -m tensorflow.python.tools.freeze_graph \
    --input_saved_model_dir=saved_model \
    --output_graph=mnist.pb \
    --output_node_names=dense_1

# 或直接转换为.tflite（更简单）
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open('mnist.tflite', 'wb') as f:
    f.write(tflite_model)
  • 节点名称确认：inputs/outputs参数必须与模型实际节点名一致。可通过!python -c "import tensorflow as tf; print([n.name for n in tf.get_default_graph().as_graph_def().node])"查看，或使用Netron工具可视化确认。
  • dataset.txt：量化校准数据集，每行一个图像文件路径（建议20~100张真实训练图像）。格式示例：
images/0.jpg
images/1.jpg
images/2.jpg
...

运行与验证

  1. 将上述转换脚本保存为convert.py，并确保mnist.pb（或.tflite）与dataset.txt位于同一目录。
  2. 执行：python convert.py
  3. 通过条件：
  • 终端输出包含build success和export success，无报错。
  • 当前目录生成mnist.rknn文件，大小约几十KB（int8量化后显著小于原始float32模型）。
  • 验证代码打印的推理输出shape为(1, 10)，与MNIST的10类输出一致。

常见问题

  • 加载.pb时提示找不到节点：使用Netron打开.pb文件，确认输入/输出张量名称后更新inputs/outputs参数。
  • 量化后精度下降：默认量化对MNIST这类简单任务影响很小（通常<1%）。若精度异常，检查dataset.txt是否使用了与训练集同分布的图像，并增加校准图像数量。
  • target_platform不匹配：根据实际部署芯片修改target_platform参数，例如RK3566/RK3588/RK3399Pro等。
← 上一章 📖 返回目录 下一章 →