当前位置: 首页 > news >正文

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版本导出

主要需要更改的位置为:

  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
  1. 将type类型进行更改(export.py)。
    # rknn模式: shape = tuple(y[0].shape)# trt模式:# shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
  1. (可选)为每个输出头分配一个名称,若不分配,会默认按照节点的名称自动进行分配。
    # 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多头输出版本,去除了解码的模块,所以在训练和推理模式需要把代码进行还原。
http://www.dtcms.com/a/273026.html

相关文章:

  • Renren框架DistributeLock排他锁实现详解
  • 企业内网系统:从传统开发到智能赋能的进化之路
  • 安达发|医疗器械行业APS自动排单:智能化生产管理的未来之路
  • useRef跨渲染周期存储
  • 数据结构 --- 队列
  • 10.Docker安装mysql
  • chatgpt是怎么诞生的,详解GPT1到GPT4的演化之路及相关背景知识
  • dexie 前端数据库封装
  • 使用快捷键迅速校准多个通道 | IPEmotion
  • 软件技术:柯里化
  • 《PyQt6-3D应用开发技术文档》
  • 仿豆包智能输入框实现
  • python基础25_某大网校(下)处理json数据以及保存题库
  • 安全访问云端内部应用:用frp的stcp功能解决SSH转发的痛点
  • Linux驱动开发(platform 设备驱动)
  • 老题新解|矩阵转置
  • AI驱动的业务系统智能化转型:从非结构化到结构化的智能转换
  • 【STM32 学习笔记】FLASH闪存
  • pytorch学习-12循环神经网络(基础篇)
  • 机器视觉之激光码检测系统
  • 【世纪龙科技】学测-汽车信息化综合实训考核平台(机电方向)
  • 数字孪生系统如何助力汽车零部件企业实现虚拟智控
  • RedisJSON 内存占用剖析与调优
  • Lua嵌入式爬虫实现步骤
  • 【Linux系统】冯诺依曼体系结构 | 初识操作系统
  • 生产者、消费者问题(C语言、POSIX)
  • 测试覆盖标准-条件覆盖-短路求值
  • 全新开源AI知识库系统!PandaWiki一键构建智能文档,支持AI问答、创作与搜索!
  • [特殊字符] 05_Jenkins 部署前端项目实现自动化部署
  • rv1106使用笔记