因工作需要,需要将目标检测模型 部署在开发板上。在走了很多弯路后 找到一个成功的案例并记载下来
这里说一下我现有的硬件设备 。
我是购买的RADXA的rock3a开发板 搭载的soc是rk3568
这是开发板的正面图,因为瑞芯微针对计算机视觉中的目标检测模型有一套自己的前向推理框架,所以我就着眼于搭载rockchip的开发板rock3a
目标检测模型 这里采用的是yolo模型 由于原生yolov5模型里面有一些 算子 可能在模型转换时不支持,这里采用瑞芯微官方推荐的yolov5,链接在下面:
GitHub - airockchip/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite
下载好源代码后,需要修改一个地方:
修改的代码我这里贴出来
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
if os.getenv('RKNN_model_hack', '0') != '0':
z.append(torch.sigmoid(self.m[i](x[i])))
continue
x[i] = self.m[i](x[i]) # conv
'''
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # inference
if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
y = x[i].sigmoid()
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy, wh, conf = y.split((2, 2, self.nc + 1), 4) # y.tensor_split((2, 4, 5), 4) # torch 1.8.0
xy = (xy * 2 + self.grid[i]) * self.stride[i] # xy
wh = (wh * 2) ** 2 * self.anchor_grid[i] # wh
y = torch.cat((xy, wh, conf), 4)
z.append(y.view(bs, -1, self.no))
if os.getenv('RKNN_model_hack', '0') != '0':
return z
return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
'''
return x[0],x[1],x[2]
然后 根据自己的数据集 训练一版模型 ,训练过程这里不赘述。
这里将训练好的模型yolov5s.pt 转化为onnx模型
转化命令为 :
python3 export.py --weights yolov5s.pt --img 640 --batch 1 --opset 12 --include onnx
转化完后 需要用Netron软件打开模型的可视化界面
记下 图中 的三个框里面的节点名称,这三个节点是网络的输出节点
新建一个 python文件 onnx2rknn.py 该文件是仿制test.py文件编写的
文件内容为
from rknn.api import RKNN
ONNX_MODEL = 'yolov5s.onnx'
platform = "rk3568"
RKNN_MODEL = 'yolov5s_{}_out_opt.rknn'.format(platform)
if __name__ == '__main__':
add_perm = False # 如果设置成True,则将模型输入layout修改成NHWC
# Create RKNN object
rknn = RKNN(verbose=True)
# pre-process config
print('--> config model')
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform,
)
print('done')
# Load tensorflow model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['output', '327', '328'])# 这里一定要根据onnx模型修改
if ret != 0:
print('Load onnx model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
print('Build rkmodel failed!')
exit(ret)
print('done')
# rknn.export_rknn_precompile_model(RKNN_MODEL)
rknn.export_rknn(RKNN_MODEL)
rknn.release()
并将该文件放置于rknn-toolkit工作目录下
我的目录是/home/rock/workspace/rknn-toolkit2-1.4.0/examples/onnx/yolov5
也就是example目录下的onnx专栏的yolov5子目录
注意 需要在onnx2rknn.py文件中 注明 模型运行的目标平台 我的目标平台是rk3568
另外需要将上一步 netron 查看的输出部分的三个节点 在onnx2rknn.py文件中注明
也就是如下截图所示
python3 onnx2rknn.py 运行该文件后 生成的rknn模型
将该模型 传输到开发板上
注意,这里有一个容易忽略的地方 我的rknn-toolkit的版本是1.4.0
如果我的开发板librknnrt.so 动态库的版本是1.3.0 就会报错 这里版本一定要对应上
在开发板上运行模型 进行单张图片的目标检测 ,运行结果如下图所示
文章来源:https://www.toymoban.com/news/detail-510624.html
文章来源地址https://www.toymoban.com/news/detail-510624.html
到了这里,关于采用rknn-toolkit导出rknn模型并部署在rock3a-rk3568芯片 上全流程的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!