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

TensorRT 7 `anchor_grid` compatibility fix (#6185) 

* fix: TensorRT 7 incompatiable

* Add comment

* Add if: else and comment

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/6195/head
imyhxy 2022-01-05 03:09:25 +08:00 committed by GitHub
parent fb839298a1
commit a2f4a1799b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
export.py
View File

@ -175,7 +175,13 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
import tensorrt as trt
opset = (12, 13)[trt.__version__[0] == '8'] # test on TensorRT 7.x and 8.x
export_onnx(model, im, file, opset, train, False, simplify)
if opset == 12: # 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, opset, train, False, simplify)
model.model[-1].anchor_grid = grid
else: # TensorRT >= 8
export_onnx(model, im, file, opset, train, False, simplify)
onnx = file.with_suffix('.onnx')
assert onnx.exists(), f'failed to export ONNX file: {onnx}'
@ -418,12 +424,12 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
# Exports
if 'torchscript' in include:
export_torchscript(model, im, file, optimize)
if 'engine' in include: # TensorRT required before ONNX
export_engine(model, im, file, train, half, simplify, workspace, verbose)
if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
export_onnx(model, im, file, opset, train, dynamic, simplify)
if 'openvino' in include:
export_openvino(model, im, file)
if 'engine' in include:
export_engine(model, im, file, train, half, simplify, workspace, verbose)
if 'coreml' in include:
export_coreml(model, im, file)