【RV126 RKNN模型转换】

最新推荐文章于 2024-05-18 16:45:21 发布

@有梦想的咸鱼@

最新推荐文章于 2024-05-18 16:45:21 发布

阅读量600

点赞数 12

分类专栏： RK 深度学习文章标签：深度学习 pytorch c++ c语言

本文链接：https://blog.csdn.net/qq_45193535/article/details/136439387

版权

RK 深度学习专栏收录该内容

1 篇文章 0 订阅

订阅专栏

本文详细介绍了如何使用rknntoolkit将TensorFlow和Pytorch模型转换为RKNN格式，包括设置模型参数、加载模型、构建和导出模型的过程，以SSDMobileNetV1为例。

摘要由CSDN通过智能技术生成

一.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