HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.0.230 TensorRT export hang fix (#7180)

Browse source 
Co-authored-by: 曾逸夫（Zeng Yifu） <41098760+Zengyf-CVer@users.noreply.github.com>
Co-authored-by: CV & LLM & AIGC er <wenvoi@163.com>
Co-authored-by: Aaron <42322215+aaronllowe@users.noreply.github.com>
Co-authored-by: crypthusiast0 <42322215+crypthusiast0@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-12-26 19:17:29 -05:00 • committed by GitHub
parent bc5b528ca7
commit 8d2a32db77
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 15 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/engine/exporter.py
View file

@ -579,6 +579,8 @@ class Exporter:
def export_engine(self, prefix=colorstr('TensorRT:')):
"""YOLOv8 TensorRT export https://developer.nvidia.com/tensorrt."""
assert self.im.device.type != 'cpu', "export running on CPU but must be on GPU, i.e. use 'device=0'"
f_onnx, _ = self.export_onnx() # run before trt import https://github.com/ultralytics/ultralytics/issues/7016
try:
import tensorrt as trt # noqa
except ImportError:
@ -587,8 +589,8 @@ class Exporter:
import tensorrt as trt # noqa
check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=7.0.0
self.args.simplify = True
f_onnx, _ = self.export_onnx()
LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
assert Path(f_onnx).exists(), f'failed to export ONNX file: {f_onnx}'