Pytorch模型（基于SSD的目标检测算法）如何转为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
net.eval()
trace_model = torch.jit.trace(net, torch.Tensor(1,3,224,224))#Tensor也需要根据自己的模型进行修改
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)#该函数对于torch<1.6的可能不支持，我的torch版本为1.6
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]]#输入模型只接受固定尺寸输入
# Create RKNN object
rknn = RKNN()
# pre-process config
print('--> config model')
rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]], reorder_channel='0 1 2')
#对于输入数据是三通道的(Cin0, Cin1, Cin2)数据来讲，经过预处理后，输出的数据为(Cout0,Cout1, Cout2)计算过程如下:
#Cout0 = (Cin0 - M0)/S0, Cout1 = (Cin1 - M1)/S0, Cout2 = (Cin2 - M2)/S0,其中mean_values=[[M0,M1,M2]]
#std_values=[[S0,S1,S2]]。一般情况使用mean_values=[[0,0,0]], std_values=[[255,255,255]]要将输入数据归一化到[0, 1]
#之间，使用mean_values=[[0,0,0]], std_values=[[128,128,128]]将输入数据归一化到[-1, 1]之间；其他情况需要你自己计算了！！！！
#reorder_channel：表示是否需要对图像通道顺序进行调整。’0 1 2’表示按照输入的通
#道顺序来推理，比如图片输入时是 RGB，那推理的时候就根据 RGB 顺序传给输入层；’2
#1 0’表示会对输入做通道转换，比如输入时通道顺序是 RGB，推理时会将其转成 BGR，
#再传给输入层，同样的，输入时通道的顺序为 BGR 的话，会被转成 RGB 后再传给输入层。
print('done')
# Load pytorch model
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')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False)
if ret != 0:
print('Build pytorch failed!')
exit(ret)
print('done')
# Export 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()
（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]]#输入模型只接受固定尺寸输入
# Create RKNN object
rknn = RKNN()
# pre-process config
print('--> config model')
rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]], reorder_channel='0 1 2')
print('done')
# Load pytorch model
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')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False)
if ret != 0:
print('Build pytorch failed!')
exit(ret)
print('done')
# Export 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()
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()
文章出处登录后可见！
已经登录？立即刷新
