TensoRT代码可参考:https://github.com/shouxieai/tensorRT_Pro
模型链接:https://pan.baidu.com/s/18yIsypWMg0sT_uAR_MhDSA 提取码: sh7c
为了正确导出onnx需要将models/yolo.py文件中56行
# x = x if self.training else (torch.cat(z, 1),x) # 修改为
x = x if self.training else (torch.cat(z, 1))
在detect.py中采用如下方式导出
dummy = torch.zeros(1,3,640,640)
torch.onnx.export(
model,(dummy,),
"yolov7.onnx",
input_names=["image"],
output_names=["predict"],
opset_version=13,
dynamic_axes={"image":{0:"batch"},"predict":{0:"batch"},}
)