下面的内容是对 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行: tr

本文介绍了YOLOv5在6.0和5.0版本中导出ONNX模型以备TensorRT部署的详细步骤。内容包括修改导出文件以简化模型,解释了模型输出结构,讨论了inplace操作的影响,以及在不同版本中如何处理网格初始化问题。此外,还提供了避免问题的建议和参考资源。
订阅专栏 解锁全文
5435

被折叠的 条评论
为什么被折叠?



