当前位置:   article > 正文

yolov5训练自己的pt文件,转onnx,再转成rknn,到RK3588开发板运行测试_训练好的算法pt文件后面怎么办

训练好的算法pt文件后面怎么办

一、训练自己的模型

  1. yolov5训练好自己的模型,例如训练完后,名称为best.pt,路径为runs/exp/weights/best.pt。

  1. 采用detect.py文件验证best.pt可以正常检测目标,再进行下一步工作。

二、pt转onnx

  1. 修改utils/yolo.py文件的后处理部分,将class Detect(nn.Module) 类的子函数forward由

def forward(self, x):        z = []  # inference outputfor i in range(self.nl):            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()ifnot self.training:  # inferenceif self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:                    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. - 0.5 + self.grid[i]) * self.stride[i]  # xy                    y[...,2:4] = (y[...,2:4] * 2) ** 2 * self.anchor_grid[i]  # whelse:  # for YOLOv5 on AWS Inferentia https://github
声明:本文内容由网友自发贡献,转载请注明出处:【wpsshop博客】
推荐阅读
相关标签
  

闽ICP备14008679号