当前位置:   article > 正文

采用rknn-toolkit导出rknn模型并部署在rock3a-rk3568芯片 上全流程_rknn killed build

rknn killed build

因工作需要,需要将目标检测模型 部署在开发板上。在走了很多弯路后 找到一个成功的案例并记载下来

这里说一下我现有的硬件设备 。

我是购买的RADXA的rock3a开发板 搭载的soc是rk3568

这是开发板的正面图,因为瑞芯微针对计算机视觉中的目标检测模型有一套自己的前向推理框架,所以我就着眼于搭载rockchip的开发板rock3a

目标检测模型 这里采用的是yolo模型  由于原生yolov5模型里面有一些 算子 可能在模型转换时不支持,这里采用瑞芯微官方推荐的yolov5,链接在下面:

GitHub - airockchip/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite

下载好源代码后,需要修改一个地方:

修改的代码我这里贴出来

  1. def forward(self, x):
  2. z = [] # inference output
  3. for i in range(self.nl):
  4. if os.getenv('RKNN_model_hack', '0') != '0':
  5. z.append(torch.sigmoid(self.m[i](x[i])))
  6. continue
  7. x[i] = self.m[i](x[i]) # conv
  8. '''
  9. bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
  10. x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
  11. if not self.training: # inference
  12. if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
  13. self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
  14. y = x[i].sigmoid()
  15. if self.inplace:
  16. y[..., 0:2] = (y[..., 0:2] * 2 + self.grid[i]) * self.stride[i] # xy
  17. y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
  18. else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
  19. xy, wh, conf = y.split((2, 2, self.nc + 1), 4) # y.tensor_split((2, 4, 5), 4) # torch 1.8.0
  20. xy = (xy * 2 + self.grid[i]) * self.stride[i] # xy
  21. wh = (wh * 2) ** 2 * self.anchor_grid[i] # wh
  22. y = torch.cat((xy, wh, conf), 4)
  23. z.append(y.view(bs, -1, self.no))
  24. if os.getenv('RKNN_model_hack', '0') != '0':
  25. return z
  26. return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
  27. '''
  28. 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文件编写的

文件内容为

  1. from rknn.api import RKNN
  2. ONNX_MODEL = 'yolov5s.onnx'
  3. platform = "rk3568"
  4. RKNN_MODEL = 'yolov5s_{}_out_opt.rknn'.format(platform)
  5. if __name__ == '__main__':
  6. add_perm = False # 如果设置成True,则将模型输入layout修改成NHWC
  7. # Create RKNN object
  8. rknn = RKNN(verbose=True)
  9. # pre-process config
  10. print('--> config model')
  11. rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform,
  12. )
  13. print('done')
  14. # Load tensorflow model
  15. print('--> Loading model')
  16. ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['output', '327', '328'])# 这里一定要根据onnx模型修改
  17. if ret != 0:
  18. print('Load onnx model failed!')
  19. exit(ret)
  20. print('done')
  21. # Build model
  22. print('--> Building model')
  23. ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
  24. if ret != 0:
  25. print('Build rkmodel failed!')
  26. exit(ret)
  27. print('done')
  28. # rknn.export_rknn_precompile_model(RKNN_MODEL)
  29. rknn.export_rknn(RKNN_MODEL)
  30. 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 就会报错 这里版本一定要对应上

在开发板上运行模型 进行单张图片的目标检测 ,运行结果如下图所示

 

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/天景科技苑/article/detail/854478
推荐阅读
相关标签
  

闽ICP备14008679号