Pytorch模型（基于SSD的目标检测算法）如何转为rknn模型
实践证明，瑞星微就是个lese！RKNN-Toolkit1.3之后是支持pytorch直接转rknn的，但是存在很多小问题，例如我的模型在转化过程中是不能量化的；转化后的模型在部署到RV1126上时也会报错等等。最终还是使用老方法，先将pytorch转为onnx，再转化为rknn
一、先新建一个python环境
先新建一个python环境，例如使用anaconda。如果使用之前的python的环境，可能会有一些python库与rknn库不兼容导致疯狂报错！！！
二、选择一个RKNN-Toolkit
选择版本大于1.3的RKNN-Toolkit，因为1.3之前的版本需要先将pytorch模型转为caffe模型的ONNX，然后再将其转为rknn；1.3之后的版本支持直接将pytorch转为rknn。
地址如下：https://github.com/rockchip-linux/rknn-toolkit/releases，我选择的是RKNN-Toolkit1.7，其下会有四个压缩文件
; 其中packages中是使用说明和example等，源代码中是RKNN-Toolkit1.7各个版本的轮子。
三、配置环境
在rknn-toolkit-1.7.0\doc\Rockchip_Quick_Start_RKNN_Toolkit_V1.7.0_CN.pdf中有各个系统的环境配置指南，如下图所示，安装它指定的pytorch或tensorflow版本（最好全部安装），pytorch版本最好>1.6
; 四、pytorch的.pth权重文件转为包含模型框架（我的是SSD模型）的.pt文件
需要根据rknn-toolkit-1.7.0\examples\pytorch\resnet18\test.py文件自己修改代码
def export_pytorch_model():
net = models.resnet18(pretrained=True)
net.eval()
trace_model = torch.jit.trace(net, torch.Tensor(1,3,224,224))
trace_model.save('./resnet18.pt')
例如我的修改为如下：
model_path = '权重文件路径（.pth）'
label_path = '数据集标签路径'
class_names = [name.strip() for name in open(label_path).readlines()]
export_pytorch_model(model_path,class_names)
def export_pytorch_model(model_path,class_names):
net = create_mobilenetv1_ssd_lite_025extra(len(class_names), width_mult=0.25 ,is_test=True)
device = torch.device("cpu")
net=net.to(device)
net.load(model_path)
net.eval()
example = torch.rand(1, 3, 300, 300)
example = example.to(device)
trace_model = torch.jit.trace(net, example)
trace_model.save('./models/ssd.pt')
五、.pt文件转为.rknn文件
（1）第一种方法
这一步只需要修改几个参数就可以，如下图所示：
if __name__ == '__main__':
model_path = '权重路径'
label_path = '标签路径'
class_names = [name.strip() for name in open(label_path).readlines()]
export_pytorch_model(model_path,class_names)
model = 'models/ssd.pt'
input_size_list = [[3,300,300]]
rknn = RKNN()
print('--> config model')
rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]], reorder_channel='0 1 2')
print('done')
print('--> Loading model')
ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
if ret != 0:
print('Load pytorch model failed!')
exit(ret)
print('done')
print('--> Building model')
ret = rknn.build(do_quantization=False)
if ret != 0:
print('Build pytorch failed!')
exit(ret)
print('done')
print('--> Export RKNN model')
ret = rknn.export_rknn('models/ssd.rknn')
if ret != 0:
print('Export ssd.rknn failed!')
exit(ret)
print('Export ssd.rknn success!')
rknn.release()
（2）第二种方法：使用rknn自带的可视化程序（非常友好！！！）
在rknn-toolkit-1.7.0\doc\Rockchip_User_Guide_RKNN_Toolkit_Visualization_V1.7.0_CN.pdf中详细介绍了该可视化程序如何使用。
在配好的环境下使用如下命令，打开可视化程序
python -m rknn.bin.visualization
选择pytorch模型，进入如下界面
; 其中的一些参数与（1）中的参数一样。至于量化类型我也不太懂，所以选择了none；Batch Size的大小默认值为 100。量化时将根据该参数决定每一批次喂的数据量，以校正量化结果。如果 dataset 中的数据量小于 100，则该参数值将自动调整为dataset 中的数据量（没啥用，这个参数设置的前提是你需要进行rknn推理，如果你只想导出rknn模型，数据集和batch size随便填一个就好了）。
完整代码如下：
import numpy as np
from rknn.api import RKNN
import torch
from vision.ssd.mobilenetv1_ssd_lite_025extra import create_mobilenetv1_ssd_lite_025extra
def export_pytorch_model(model_path,class_names):
net = create_mobilenetv1_ssd_lite_025extra(len(class_names), width_mult=0.25 ,is_test=True)
device = torch.device("cpu")
net=net.to(device)
net.load(model_path)
net.eval()
example = torch.rand(1, 3, 300, 300)
example = example.to(device)
trace_model = torch.jit.trace(net, example)
trace_model.save('./models/ssd.pt')
if __name__ == '__main__':
model_path = 'models/ws0.25tmax400extra0.25/mb1-ssd-lite-025extra-Epoch-399-Loss-3.2293308803013394.pth'
label_path = 'models/voc-model-labels.txt'
class_names = [name.strip() for name in open(label_path).readlines()]
export_pytorch_model(model_path,class_names)
model = 'models/ssd.pt'
input_size_list = [[3,300,300]]
rknn = RKNN()
print('--> config model')
rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]], reorder_channel='0 1 2')
print('done')
print('--> Loading model')
ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
if ret != 0:
print('Load pytorch model failed!')
exit(ret)
print('done')
print('--> Building model')
ret = rknn.build(do_quantization=False)
if ret != 0:
print('Build pytorch failed!')
exit(ret)
print('done')
print('--> Export RKNN model')
ret = rknn.export_rknn('models/ssd.rknn')
if ret != 0:
print('Export ssd.rknn failed!')
exit(ret)
print('Export ssd.rknn success!')
rknn.release()
xport rknn model
print(‘–> Export RKNN model’)
ret = rknn.export_rknn(‘models/ssd.rknn’)
if ret != 0:
print(‘Export ssd.rknn failed!’)
exit(ret)
print(‘Export ssd.rknn success!’)
rknn.release()
Original: https://blog.csdn.net/weixin_41166529/article/details/125150024
Author: Luxus1001
Title: Pytorch模型如何转rknn模型（基于SSD的目标检测算法）
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/682857/
转载文章受原作者版权保护。转载请注明原作者出处！