[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
This commit is contained in:
A. Unique TensorFlower 2021-02-10 12:53:49 -08:00 committed by TensorFlower Gardener
parent f9529da88c
commit 452fc5dead
2 changed files with 17 additions and 3 deletions

View File

@ -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))

View File

@ -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))