2. TensorFlow模型RKNN量化转换
任务目标
将训练好的TensorFlow模型（.pb或SavedModel格式）通过RKNN-Toolkit2转换为RKNN格式，并执行默认的int8量化，最终得到可在瑞芯微NPU上部署的.rknn文件。
环境依赖
- RKNN-Toolkit2（版本≥1.6.0），安装方式：
pip install rknn-toolkit2 - TensorFlow 1.x（用于加载.pb模型）或TensorFlow 2.x（用于加载SavedModel）
- Python 3.6~3.10
完整实现代码
from rknn.api import RKNN
# 1. 创建RKNN对象
rknn = RKNN(verbose=True)
# 2. 配置量化参数（默认int8量化）
rknn.config(
mean_values=[[123.675, 116.28, 103.53]], # 与训练时预处理一致
std_values=[[58.395, 57.12, 57.375]],
target_platform='rk3588', # 根据实际NPU平台选择
quantized_dtype='int8' # 默认即为int8
)
# 3. 加载TensorFlow模型
# 方式A：加载.pb冻结图
ret = rknn.load_tensorflow(
tf_pb='./model/frozen_graph.pb',
inputs=['input_tensor'], # 输入节点名称
outputs=['softmax_tensor'], # 输出节点名称
input_size_list=[[1, 224, 224, 3]] # 输入shape
)
# 方式B：加载SavedModel（若使用方式A则注释掉B）
# ret = rknn.load_tensorflow(
# tf_path='./saved_model_dir',
# inputs=['serving_default_input_tensor'],
# outputs=['StatefulPartitionedCall:0'],
# input_size_list=[[1, 224, 224, 3]]
# )
if ret != 0:
print('模型加载失败')
exit(ret)
# 4. 构建RKNN模型（执行int8量化）
ret = rknn.build(do_quantization=True) # True表示执行量化
if ret != 0:
print('模型构建失败')
exit(ret)
# 5. 导出RKNN模型
ret = rknn.export_rknn('./model/model_int8.rknn')
if ret != 0:
print('模型导出失败')
exit(ret)
# 6. 释放资源
rknn.release()
print('转换完成，已生成 model_int8.rknn')运行方式
- 将上述代码保存为
convert_tf2rknn.py，并确保模型文件路径正确； - 执行命令：
python convert_tf2rknn.py； - 转换过程中终端会打印各阶段日志，包括量化统计信息。
验证方法
- 文件检查：确认
model_int8.rknn文件已生成且大小合理（通常为原始FP32模型的1/4左右）； - 加载验证：使用以下代码加载生成的RKNN模型，确认无报错：
from rknn.api import RKNN
rknn = RKNN()
ret = rknn.load_rknn('./model/model_int8.rknn')
if ret == 0:
print('RKNN模型加载成功')
rknn.release()
else:
print('RKNN模型加载失败')- 推理验证（可选）：准备一张与训练集同分布的测试图片，执行推理并观察输出类别是否合理。若输出与原始TensorFlow模型一致或接近，则量化转换成功。
常见问题
- 节点名称错误：使用
tensorboard或graph.pbtxt确认输入/输出节点名称，名称必须与图中完全一致； - 量化精度下降：若量化后精度明显下降，可尝试提供
dataset.txt（包含100~500张代表性图片路径）作为量化校准集，在build前调用rknn.config(quantized_algorithm='normal')并传入rknn.build(do_quantization=True, dataset='./dataset.txt')； - 平台不匹配：
target_platform需与目标NPU芯片一致（如rk3588、rk3566、rv1126等），否则NPU无法加载。