zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

TensorRT 7 export fix (#6235)

This commit is contained in:
Glenn Jocher 2022-01-07 09:31:17 -08:00 committed by GitHub
parent 33a67b4918
commit 6865d19a92
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
export.py
View File

@ -174,7 +174,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
check_requirements(('tensorrt',)) check_requirements(('tensorrt',))
import tensorrt as trt import tensorrt as trt
if trt.__version__[0] == 7: # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
grid = model.model[-1].anchor_grid grid = model.model[-1].anchor_grid
model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
export_onnx(model, im, file, 12, train, False, simplify) # opset 12 export_onnx(model, im, file, 12, train, False, simplify) # opset 12