【目标检测-YOLO】YOLOv5-yolov5s TensorRT部署准备之ONNX导出(第一篇)

下面的内容是对 YOLOv5-yolov5s TensorRT部署前的准备之导出可用的ONNX。

之前已经写过部分内容,但是还不够详细。

1. YOLOv5-6.0版本

以下内容以 6.0 版本为准。

修改导出文件:

    def forward(self, x):
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv2d
            # 修改1:对于任何用到shape、size返回值的参数时,为了避免pytorch 导出 onnx 时候,对size 进行跟踪,跟踪时候会生成gather、shape的节点。使用 map 加上int转换。
            bs, _, ny, nx = map(int,x[i].shape)  # x(bs,255,20,20) to x(bs,3,20,20,85)
            # 修改2:对于reshape、view操作时,-1的指定请放到batch维度。其他维度可以计算出来即可。batch维度禁止指定为大于-1的明确数字。所以这里 bs 改为 -1.

            x[i] = x[i].view(-1, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            # export.py:model.train() if train else model.eval()  # 264行: train=False,  # model.train() mode 所以这里是 model.eval()模式, self.training 为 False,所以推理使用以下代码
            if not sel

Original: https://blog.csdn.net/hymn1993/article/details/123893719
Author: 理心炼丹
Title: 【目标检测-YOLO】YOLOv5-yolov5s TensorRT部署准备之ONNX导出(第一篇)

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/629356/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球