【无标题】

主要为了tensorrt的部署加速

课程来自linkhttps://ke.qq.com/course/4993141?quicklink=1#term_id=105165326
这里使用的是[yolov5-6.0 地址linkhttps://github.com/ultralytics/yolov5/tree/v6.0
yolov5的训练就不细讲了，这里假设你已经训练好了自己的模型

1.修改export.py

只将batch的维度设为动态，其他维度固定，
在第73行修改为

torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
                          training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                          do_constant_folding=not train,
                          input_names=['images'],
                          output_names=['output'],
                          dynamic_axes={'images': {0: 'batch'},  # shape(1,3,640,640)
                                        'output': {0: 'batch'}  # shape(1,25200,85)
                                        } if dynamic else None)

2.修改model文件下yolo.py

36行找到Detect类，没有 control f 搜索detect ，
将56行中

bs, _, ny, nx = x[i].shape

修改为

bs, _, ny, nx = map(int,x[i].shape)

修改72行 ，原因和第一条一致只降batch维度设为动态

z.append(y.view(bs, self.na*ny*nx, self.no))

修改74行 ruturn 主要是原来yolo模型输出为4 ，但是只需要一个输出

return x if self.training else torch.cat(z, 1)

最后修改63行 ，原因是作者原来的 anchor_grid 会生成一堆节点

anchor_grid = (self.anchors[i].clone() * self.stride[i]).view(1,-1,1,1,2)

最后导出命令为
python export.py --weights=./runs/train/exp2/weights/last.pt --dynamic --include=onnx --opset=11