一.RKNN模型转换流程
二.模型转换为RKNN模型
1.tensorflow模型转为RKNN
(1)使用vscode调用rknntoolkit提供的接口编写转换代码
以tensorflow模型为例
from rknn.api import RKNN
if name == 'main':
确定目标设备target
target = 'rv1126'# 创建RKNN对象
rknn = RKNN()
#配置RKNN模型
print('--> config model')
rknn.config(mean_values=[[127.5, 127.5, 127.5]],
std_values=[[127.5, 127.5, 127.5]],
reorder_channel='0 1 2',
target_platform=[target])
print('done')
#加载TF模型
print('--> loading model')
ret = rknn.load_tensorflow(tf_pb='./original_model/ssd_mobilenet_v1_coco_2018_01_28/frozen_inference_graph.pb',
inputs=['FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_0/BatchNorm/batchnorm/mul_1'],
outputs=['concat', 'concat_1'],
input_size_list=[[300, 300, 3]])
if ret != 0:
print('load model failed!')
rknn.release()
exit(ret)
print('done')
构建 RKNN 模型
print('--> building model')
ret = rknn.build(do_quantization=True, dataset='./dataset/dataset3.txt')
if ret != 0:
print('build model failed!')
rknn.release()
exit(ret)
print('done')
#导出 RKNN 模型
print('--> export RKNN model')
rknn.export_rknn('./rknn_model/ssd_mobilenet_v1_coco.rknn')
if ret != 0:
print('export RKNN model failed!')
rknn.release()
exit(ret)
print('done')
释放RKNN对象
rknn.release()
(2)执行运行转换的python代码(先清除说有输出然后全部运行)
2.Pytorch模型转换为RKNN
Pytorch转换TorchScript 或者ONNX 格式的文件的操作步骤
(1) 按照链接https://github.com/ultralytics/yolov5/issues/251 的指导,在训练的虚拟环境下(需要有训
练的GPU,不要在VMware下安装的Ubuntu中操作,否则会报错找不到libcufft.so.11)安装依赖的环境:
// 原来的文件默认安装torchvision>=0.8.1,装完后可实现将.pt文件转换为.torchvision格式文件
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
// 若需要将.pt文件转换为.onnx格式文件,仍需要再安装onxx,推荐onnx>=1.9.0
pip install onnx==1.12.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
(2)执行如下命令,将.pt文件导出.onnx或者. torchscript格式文件
// 转换为TorchScript格式,得到best.torchscript文件
python export.py --weights ./runs/train/exp/weights/best.pt --img 640 --batch 1 --include torchscript
// 转换为ONNX格式,得到best.onnx文件
python export.py --weights ./runs/train/exp/weights/best.pt --img 640 --batch 1 --include onnx
// 若需要指定opset,可根据安装的onnx库的版本来调整,如安装的onnx库的版本是1.12.0,后面加上--opset 12
python export.py --weights ./runs/train/exp/weights/best.pt --img 640 --batch 1 --include onnx --opset 12
三.模型文件目录介绍
以为tensorflow模型架构下的ssd_mobilenet_v1_coco 为例