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 output
for i in range(self.nl):
x[i] = torch.sigmoid(self.m[i](x[i])) # conv
return 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多头输出版本，去除了解码的模块，所以在训练和推理模式需要把代码进行还原。