官方预训练模型转换
- 下载yolov5-v6.0分支源码解压到本地,并配置基础运行环境。
- 下载官方预训练模型
- 进入yolov5-6.0目录下,新建文件夹weights,并将步骤2中下载的权重文件放进去。
- 修改models/yolo.py文件
def forward(self, x):
z = [] # inference output
for 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()
if not self.training: # inference
if 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] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
return x if self.training else torch.cat(z, 1)
- 命令行执行
python3 export.py --weights weights/yolov5n.pt --include onnx --opset 11 --simplify
(默认为yolov5n.pt,转换成功会输出一下信息, 转换后的模型存于权重同级目录yolov5n.onnx后缀模型) 注意opset仅支持10/11
ONNX: simplifying with onnx-simplifier 0.4.33...
ONNX: export success, saved as weights/yolov5n.onnx (7.9 MB)
ONNX: run --dynamic ONNX model inference with: 'python detect.py --weights weights/yolov5n.onnx'
Export complete (1.31s)
地平线BPU加载推理篇
- 使用地平线BPU SDK接口加载推理模型, 后续会放出代码。