网站开发的教学网站金华百度推广公司

YOLO系列pt导出不同onnx方法

1 YOLOv5模型导出

1.1 默认版本导出

直接使用export.py进行导出非dynamic(动态batch)的onnx模型，输出维度为(batch，25200,85)。其中batch为执行export.py时指定；25200为(8080+4040+20*20)*3；85为(x,y,w,h,obj_conf,cls1_conf,cls2_conf…)。

1.2 RKNN版本导出

主要需要更改的位置为：

1. 取消对于目标框的解码模块，仅获取对应的3个检测头即可(yolo.py)。

    def forward(self, x):"""Processes input through YOLOv5 layers, altering shape for detection: `x(bs, 3, ny, nx, 85)`."""# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> 测试导出符合trt的onnx模型 <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<# z = []  # inference output# for i in range(self.nl):#     # self.m:针对最后3个输出层之前的卷积1*1conv(ModuleList:1*1 conv)#     x[i] = self.m[i](x[i])#     bs, _, ny, nx = x[i].shape  # x(bs,255,20,20)#     print(f"bs:{bs}, ny:{ny}, nx:{nx}")#     x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() # x(bs,3,20,20,85), 分开anchor方便按grid进行遍历#     # 在执行export.py时, self.export为False, self.training为False#     if not self.training:#         # 创建网格(grid offset)和anchor尺度#         if self.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)#         if isinstance(self, Segment):  # (boxes + masks)#             xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)#             xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i]  # xy#             wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i]  # wh#             y = torch.cat((xy, wh, conf.sigmoid(), mask), 4)#         else:  # Detect (boxes only)#             xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4) # 按照最后一维(dim=4)进行拆分: xy([bs, 3, 20, 20, 2]), wh([bs, 3, 20, 20, 2]), conf([bs, 3, 20, 20, 81])#             xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy:相较于当前cell的偏移量; self.grid[i]):当前feature下每个cell左上角坐标, self.stride[i]:映射回原图#             wh = (wh * 2) ** 2 * self.anchor_grid[i]  # 尺度缩放: self.anchor_grid[i]:当前feature下每个anchor的模板#             y = torch.cat((xy, wh, conf), 4)   # 再组装回(bs,3,20,20,85)-(x,y,w,h,obj_conf,cls_conf1...)#         z.append(y.view(bs, self.na * nx * ny, self.no)) # y:[bs, 3, ny, nx, 85] → view 成 [bs, total_boxes, no] 保存到 z, 🎯这也是trt模型的原因# '''# if self.training:#     return x# else:#     if self.export:#         return (torch.cat(z, 1),)  # 导出模式：返回元组(预测输出,)#     else:#         return (torch.cat(z, 1), x)  # 普通推理模式：返回(预测, 原始特征层输出)# '''# return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> 测试导出符合rknn的onnx模型 <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<z = []  # inference outputfor i in range(self.nl):x[i] = torch.sigmoid(self.m[i](x[i]))  # convreturn x

2. 将type类型进行更改(export.py)。

    # rknn模式: shape = tuple(y[0].shape)# trt模式:# shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape

3. (可选)为每个输出头分配一个名称，若不分配，会默认按照节点的名称自动进行分配。

    # torch.onnx.export(#     model.cpu() if dynamic else model,  # --dynamic only compatible with cpu#     im.cpu() if dynamic else im,#     f,#     verbose=False,#     opset_version=opset,#     do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False#     input_names=["images"],#     output_names=output_names,#     dynamic_axes=dynamic or None,   # )# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> rknn <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<torch.onnx.export(model.cpu(), im.cpu(), f, verbose=False, opset_version=opset,do_constant_folding=True,input_names=['image'],output_names=['output1', 'output2', 'output3'],dynamic_axes={name: {0: "B"} for name in ['image'] + ['output1', 'output2', 'output3']} if dynamic else None)#   dynamic_axes={name: {0: "B"} for name in ['data'] + ['output1', 'output2', 'output3', 'output4']} if dynamic else None)

tips:

1. 在进行导出时self.training和self.export均为false。
2. 由于导出RKNN多头输出版本，去除了解码的模块，所以在训练和推理模式需要把代码进行还原。