From 5c09aa863c866dbdb68b979a37de1936b8cc1f72 Mon Sep 17 00:00:00 2001 From: stevenlix Date: Thu, 14 Mar 2019 22:25:16 -0700 Subject: [PATCH] python binding test updates for TensorRT (#635) * Disable unsupported python tests for TensorRT * Rename python wheel for TensorRT --- setup.py | 5 ++++- tools/ci_build/build.py | 19 ++++++++++++------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/setup.py b/setup.py index 5de6018568..7b302e6a48 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,10 @@ import platform import sys package_name = 'onnxruntime' -if '--use_cuda' in sys.argv: +if '--use_tensorrt' in sys.argv: + package_name = 'onnxruntime-gpu-tensorrt' + sys.argv.remove('--use_tensorrt') +elif '--use_cuda' in sys.argv: package_name = 'onnxruntime-gpu' sys.argv.remove('--use_cuda') diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 8d3f2c24a1..64c7950d6d 100755 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -472,7 +472,7 @@ def setup_tensorrt_vars(args): return tensorrt_home -def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs, enable_python_tests, enable_tvm = False): +def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs, enable_python_tests, enable_tvm = False, enable_tensorrt = False): for config in configs: log.info("Running tests for %s configuration", config) cwd = get_config_build_dir(build_dir, config) @@ -481,6 +481,9 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs, enab cwd=cwd, dll_path=dll_path) if enable_python_tests: + # Disable python tests for TensorRT because many tests are not supported yet + if enable_tensorrt: + return if is_windows(): cwd = os.path.join(cwd, config) run_subprocess([sys.executable, 'onnxruntime_test_python.py'], cwd=cwd, dll_path=dll_path) @@ -543,12 +546,14 @@ def run_onnx_tests(build_dir, configs, onnx_test_data_dir, provider, enable_para else: run_subprocess([exe] + cmd, cwd=cwd) -def build_python_wheel(source_dir, build_dir, configs, use_cuda): +def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_tensorrt): for config in configs: cwd = get_config_build_dir(build_dir, config) if is_windows(): cwd = os.path.join(cwd, config) - if use_cuda: + if use_tensorrt: + run_subprocess([sys.executable, os.path.join(source_dir, 'setup.py'), 'bdist_wheel', '--use_tensorrt'], cwd=cwd) + elif use_cuda: run_subprocess([sys.executable, os.path.join(source_dir, 'setup.py'), 'bdist_wheel', '--use_cuda'], cwd=cwd) else: run_subprocess([sys.executable, os.path.join(source_dir, 'setup.py'), 'bdist_wheel'], cwd=cwd) @@ -685,7 +690,7 @@ def main(): build_targets(cmake_path, build_dir, configs, args.parallel) if (args.test): - run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs, args.enable_pybind, args.use_tvm) + run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs, args.enable_pybind, args.use_tvm, args.use_tensorrt) # run the onnx model tests if requested explicitly. if (args.enable_onnx_tests): # directory from ONNX submodule with ONNX test data @@ -693,8 +698,8 @@ def main(): if is_windows() or not os.path.exists(onnx_test_data_dir): onnx_test_data_dir = os.path.join(source_dir, "cmake", "external", "onnx", "onnx", "backend", "test", "data") if args.use_tensorrt: - onnx_test_data_dir = '' - run_onnx_tests(build_dir, configs, onnx_test_data_dir, 'tensorrt', False, 1) + # Disable onnx unit tests for TensorRT because many tests are not supported yet + run_onnx_tests(build_dir, configs, '', 'tensorrt', False, 1) elif args.use_cuda: run_onnx_tests(build_dir, configs, onnx_test_data_dir, 'cuda', False, 2) elif args.x86 or platform.system() == 'Darwin': @@ -708,7 +713,7 @@ def main(): if args.build: if args.build_wheel: - build_python_wheel(source_dir, build_dir, configs, args.use_cuda) + build_python_wheel(source_dir, build_dir, configs, args.use_cuda, args.use_tensorrt) log.info("Build complete")