From d93d9080e6adf002bbbc57aa2667ffefb68764b5 Mon Sep 17 00:00:00 2001 From: Jules Laplace Date: Fri, 30 Nov 2018 22:20:58 +0100 Subject: trt5 legacy api --- run_engine.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/run_engine.py b/run_engine.py index 700494d..6ac2ec9 100644 --- a/run_engine.py +++ b/run_engine.py @@ -2,7 +2,11 @@ import os import sys from random import randint import numpy as np -import tensorrt + +if 'TensorRT-5.0' in os.getenv('LD_LIBRARY_PATH'): + import tensorrt.legacy +else: + import tensorrt try: from PIL import Image @@ -19,9 +23,14 @@ pip(3) install tensorrt[examples] exit(1) try: - import tensorrt as trt - from tensorrt.parsers import caffeparser - from tensorrt.parsers import onnxparser + if 'TensorRT-5.0' in os.getenv('LD_LIBRARY_PATH'): + import tensorrt.legacy as trt + from tensorrt.legacy.parsers import caffeparser + from tensorrt.legacy.parsers import onnxparser + else: + import tensorrt as trt + from tensorrt.parsers import caffeparser + from tensorrt.parsers import onnxparser except ImportError as err: sys.stderr.write("""ERROR: failed to import module ({}) Please make sure you have the TensorRT Library installed -- cgit v1.2.3-70-g09d2