From 3eb0fb6c0399a561cf6e3c5b861cc1c268c02ed3 Mon Sep 17 00:00:00 2001
From: cjh9368 <cjh9368@163.com>
Date: Sat, 22 Aug 2020 14:10:04 +0800
Subject: [PATCH] benchmark device type check

---
 mindspore/lite/tools/benchmark/benchmark.cc | 14 +++++++++++---
 mindspore/lite/tools/benchmark/benchmark.h  |  2 +-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/mindspore/lite/tools/benchmark/benchmark.cc b/mindspore/lite/tools/benchmark/benchmark.cc
index 5fdb5e8b4..ee49d780a 100644
--- a/mindspore/lite/tools/benchmark/benchmark.cc
+++ b/mindspore/lite/tools/benchmark/benchmark.cc
@@ -542,6 +542,11 @@ int Benchmark::Init() {
     return RET_ERROR;
   }
 
+  if (_flags->device != "CPU" && _flags->device != "GPU") {
+    MS_LOG(ERROR) << "Device type:" << _flags->device << " is not supported.";
+    return RET_ERROR;
+  }
+
   return RET_OK;
 }
 
@@ -574,10 +579,13 @@ int RunBenchmark(int argc, const char **argv) {
     return RET_ERROR;
   }
 
-  if (flags.device == "NPU") {
-    status = mBenchmark.RunBenchmark("NPU");
-  } else {
+  if (flags.device == "GPU") {
+    status = mBenchmark.RunBenchmark("GPU");
+  } else if (flags.device == "CPU") {
     status = mBenchmark.RunBenchmark("CPU");
+  } else {
+    MS_LOG(ERROR) << "Device type" << flags.device << " not support.";
+    return RET_ERROR;
   }
 
   if (status != 0) {
diff --git a/mindspore/lite/tools/benchmark/benchmark.h b/mindspore/lite/tools/benchmark/benchmark.h
index 7aeaee559..b01b29b74 100644
--- a/mindspore/lite/tools/benchmark/benchmark.h
+++ b/mindspore/lite/tools/benchmark/benchmark.h
@@ -56,7 +56,7 @@ class MS_API BenchmarkFlags : public virtual FlagParser {
     AddFlag(&BenchmarkFlags::inDataPath, "inDataPath", "Input data path, if not set, use random input", "");
     AddFlag(&BenchmarkFlags::inDataTypeIn, "inDataType", "Input data type. img | bin", "bin");
     AddFlag(&BenchmarkFlags::omModelPath, "omModelPath", "OM model path, only required when device is NPU", "");
-    AddFlag(&BenchmarkFlags::device, "device", "CPU | NPU | GPU", "CPU");
+    AddFlag(&BenchmarkFlags::device, "device", "CPU | GPU", "CPU");
     AddFlag(&BenchmarkFlags::cpuBindMode, "cpuBindMode",
             "Input -1 for MID_CPU, 1 for HIGHER_CPU, 0 for NO_BIND, defalut value: 1", 1);
     // MarkPerformance
-- 
GitLab