mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

TensorRT 7 export fix ()

pull/6243/head
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

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',))
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
model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
export_onnx(model, im, file, 12, train, False, simplify) # opset 12