YOLO系列pt导出不同onnx方法
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版本导出
主要需要更改的位置为：
- 取消对于目标框的解码模块，仅获取对应的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
- 将type类型进行更改(export.py)。
# rknn模式: shape = tuple(y[0].shape)# trt模式:# shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape
- (可选)为每个输出头分配一个名称，若不分配，会默认按照节点的名称自动进行分配。
# 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:
- 在进行导出时self.training和self.export均为false。
- 由于导出RKNN多头输出版本，去除了解码的模块，所以在训练和推理模式需要把代码进行还原。