From 59a350346fd3a55125f2cf7746e69f30a9350420 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 22 Apr 2022 10:41:12 -0700 Subject: [PATCH 1/6] Autoinstall TensorRT if missing May resolve https://github.com/ultralytics/yolov5/issues/7464 --- export.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/export.py b/export.py index 2a5eff23c1a6..ead614319974 100644 --- a/export.py +++ b/export.py @@ -84,7 +84,7 @@ def export_formats(): ['TensorFlow GraphDef', 'pb', '.pb', True], ['TensorFlow Lite', 'tflite', '.tflite', False], ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False], - ['TensorFlow.js', 'tfjs', '_web_model', False],] + ['TensorFlow.js', 'tfjs', '_web_model', False], ] return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'GPU']) @@ -217,7 +217,14 @@ def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')): def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt try: - import tensorrt as trt # pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com + assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' + try: + import tensorrt as trt + except: + s = f"{prefix} tensorrt not found and is required by YOLOv5" + LOGGER.info(f"{s}, attempting auto-update...") + r = '-U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com' + LOGGER.info(subprocess.check_output(f"pip install '{r}'", shell=True).decode()) if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 grid = model.model[-1].anchor_grid @@ -230,7 +237,6 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F onnx = file.with_suffix('.onnx') LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') - assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' assert onnx.exists(), f'failed to export ONNX file: {onnx}' f = file.with_suffix('.engine') # TensorRT engine file logger = trt.Logger(trt.Logger.INFO) @@ -428,9 +434,9 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, ' - r'"Identity_1": {"name": "Identity_1"}, ' - r'"Identity_2": {"name": "Identity_2"}, ' - r'"Identity_3": {"name": "Identity_3"}}}', json) + r'"Identity_1": {"name": "Identity_1"}, ' + r'"Identity_2": {"name": "Identity_2"}, ' + r'"Identity_3": {"name": "Identity_3"}}}', json) j.write(subst) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') From f0af29c8dc30dbc41371a85039bbf296edbb267a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 Apr 2022 17:42:01 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- export.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/export.py b/export.py index ead614319974..668a21b5126b 100644 --- a/export.py +++ b/export.py @@ -84,7 +84,7 @@ def export_formats(): ['TensorFlow GraphDef', 'pb', '.pb', True], ['TensorFlow Lite', 'tflite', '.tflite', False], ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False], - ['TensorFlow.js', 'tfjs', '_web_model', False], ] + ['TensorFlow.js', 'tfjs', '_web_model', False],] return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'GPU']) @@ -434,9 +434,9 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, ' - r'"Identity_1": {"name": "Identity_1"}, ' - r'"Identity_2": {"name": "Identity_2"}, ' - r'"Identity_3": {"name": "Identity_3"}}}', json) + r'"Identity_1": {"name": "Identity_1"}, ' + r'"Identity_2": {"name": "Identity_2"}, ' + r'"Identity_3": {"name": "Identity_3"}}}', json) j.write(subst) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') From 95ebe902b5c4c8ee062ab1609d45a6fc7a59b8d0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 22 Apr 2022 10:48:44 -0700 Subject: [PATCH 3/6] Update export.py --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index 668a21b5126b..d10de6a1abe1 100644 --- a/export.py +++ b/export.py @@ -220,7 +220,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' try: import tensorrt as trt - except: + except Exception: s = f"{prefix} tensorrt not found and is required by YOLOv5" LOGGER.info(f"{s}, attempting auto-update...") r = '-U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com' From 20f5efe8cfacbdf6e1eda863c9fa9829d8dfa78b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 22 Apr 2022 12:16:50 -0700 Subject: [PATCH 4/6] Update export.py --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index d10de6a1abe1..b38d24fc4232 100644 --- a/export.py +++ b/export.py @@ -224,7 +224,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F s = f"{prefix} tensorrt not found and is required by YOLOv5" LOGGER.info(f"{s}, attempting auto-update...") r = '-U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com' - LOGGER.info(subprocess.check_output(f"pip install '{r}'", shell=True).decode()) + LOGGER.info(subprocess.check_output(f"pip install {r}", shell=True).decode()) if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 grid = model.model[-1].anchor_grid From df7da6894819cef39108dcc59ca112ec15deca47 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 22 Apr 2022 12:20:58 -0700 Subject: [PATCH 5/6] Update export.py --- export.py | 1 + 1 file changed, 1 insertion(+) diff --git a/export.py b/export.py index b38d24fc4232..f90518a3d67d 100644 --- a/export.py +++ b/export.py @@ -225,6 +225,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F LOGGER.info(f"{s}, attempting auto-update...") r = '-U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com' LOGGER.info(subprocess.check_output(f"pip install {r}", shell=True).decode()) + import tensorrt as trt if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 grid = model.model[-1].anchor_grid From c5777b9af0b0b01460d6688628445d6c8eb8b440 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 22 Apr 2022 12:23:32 -0700 Subject: [PATCH 6/6] Update export.py --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index f90518a3d67d..93d98c801d02 100644 --- a/export.py +++ b/export.py @@ -221,7 +221,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F try: import tensorrt as trt except Exception: - s = f"{prefix} tensorrt not found and is required by YOLOv5" + s = f"\n{prefix} tensorrt not found and is required by YOLOv5" LOGGER.info(f"{s}, attempting auto-update...") r = '-U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com' LOGGER.info(subprocess.check_output(f"pip install {r}", shell=True).decode())