From 452fc5dead9d19ab29448b92fb3586ee4a2647b6 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 10 Feb 2021 12:53:49 -0800 Subject: [PATCH] [TF:TRT] Allow GPU only model to run in model testing. Some models are optimized for GPU and can only run on GPU due to the limitation of CPU kernel supports. For example, MaxPooling2D doesn't support NCHW layout on CPU. If a model contains MaxPooling2D with NCHW, then it cannot run on CPU. PiperOrigin-RevId: 356805410 Change-Id: If9c06a5a4d0e1c854a5244ef7555d5615b179728 --- .../compiler/tensorrt/model_tests/model_handler.py | 14 +++++++++++--- .../tensorrt/model_tests/result_analyzer.py | 6 ++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/compiler/tensorrt/model_tests/model_handler.py b/tensorflow/python/compiler/tensorrt/model_tests/model_handler.py index 3de25cd0423..9e2526d7771 100644 --- a/tensorflow/python/compiler/tensorrt/model_tests/model_handler.py +++ b/tensorflow/python/compiler/tensorrt/model_tests/model_handler.py @@ -200,8 +200,10 @@ class TestResultCollection( @property def results(self) -> Iterable[TestResult]: - return itertools.chain([self.cpu_base_result, self.gpu_base_result], - self.trt_results) + return filter( + lambda x: x is not None, + itertools.chain([self.cpu_base_result, self.gpu_base_result], + self.trt_results)) class _ModelHandlerBase(metaclass=abc.ABCMeta): @@ -629,7 +631,13 @@ class _ModelHandlerManagerBase(metaclass=abc.ABCMeta): return model.run(inputs, warmup_iterations, benchmark_iterations, **kwargs) - cpu_base_result = run_model(self._ori_model, enable_gpu=False) + # Some models include operations that can only run on GPU. + try: + cpu_base_result = run_model(self._ori_model, enable_gpu=False) + except RuntimeError as err: + logging.info("%s cannot run on CPU. Reason: %s.", + self._ori_model.model_config, err) + cpu_base_result = None gpu_base_result = run_model(self._ori_model, enable_gpu=True) trt_results = list(map(run_model, self._trt_models)) diff --git a/tensorflow/python/compiler/tensorrt/model_tests/result_analyzer.py b/tensorflow/python/compiler/tensorrt/model_tests/result_analyzer.py index ed6eb5d6015..dce4036e63d 100644 --- a/tensorflow/python/compiler/tensorrt/model_tests/result_analyzer.py +++ b/tensorflow/python/compiler/tensorrt/model_tests/result_analyzer.py @@ -125,6 +125,9 @@ def analyze_test_latency(test_results: model_handler.TestResultCollection, base_result = ( test_results.cpu_base_result if use_cpu_baseline else test_results.gpu_base_result) + if base_result is None: + raise ValueError( + f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!") base_mean_time = np.asscalar(np.mean(base_result.model_latency)) column_names = ["time(ms)", "speedup"] rows = [] @@ -148,6 +151,9 @@ def analyze_test_numerics(test_results: model_handler.TestResultCollection, base_result = ( test_results.cpu_base_result if use_cpu_baseline else test_results.gpu_base_result) + if base_result is None: + raise ValueError( + f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!") for fn0, fn1 in itertools.product(preprocess_funcs, postprocess_funcs): func0, func1 = preprocess_funcs[fn0], postprocess_funcs[fn1] column_names.append("{}_{}".format(fn0, fn1))