tflite_benchmark: use get_env_int for THREAD (#1840)
authorVitaliy Cherepanov/AI Tools Lab /SRR/Engineer/삼성전자 <v.cherepanov@samsung.com>
Tue, 3 Jul 2018 10:40:08 +0000 (13:40 +0300)
committer박세희/동작제어Lab(SR)/Principal Engineer/삼성전자 <saehie.park@samsung.com>
Tue, 3 Jul 2018 10:40:08 +0000 (19:40 +0900)
Replace std::getenv by nnfw::util::get_env_int for THREAD

Signed-off-by: Vitaliy Cherepanov <v.cherepanov@samsung.com>
tools/tflite_benchmark/src/tflite_benchmark.cc

index 9e768fe..f20cef3 100644 (file)
 
 #include <iostream>
 
+#include "util/environment.h"
 #include "util/benchmark.h"
 
 using namespace tflite;
 using namespace tflite::ops::builtin;
 
-class Thread
-{
-public:
-  // Use threads as much as possible
-  Thread() : _value(-1)
-  {
-    auto env = std::getenv("THREAD");
-
-    if (env)
-    {
-      _value = atoi(env);
-    }
-  }
-
-public:
-  int32_t value(void) const { return _value; }
-
-private:
-  int32_t _value;
-};
-
 class Count
 {
 public:
@@ -100,7 +80,6 @@ bool checkParams(const int argc, char **argv)
 
 int main(const int argc, char **argv)
 {
-  const Thread thread{};
 
   if (!checkParams(argc, argv))
   {
@@ -110,13 +89,14 @@ int main(const int argc, char **argv)
   const auto filename = argv[1];
 
   bool use_nnapi = false;
+  const auto thread_count = nnfw::util::get_env_int("THREAD", -1);
 
   if (std::getenv("USE_NNAPI") != nullptr)
   {
     use_nnapi = true;
   }
 
-  std::cout << "Num threads: " << thread.value() << std::endl;
+  std::cout << "Num threads: " << thread_count << std::endl;
   if (use_nnapi)
   {
     std::cout << "Use NNAPI" << std::endl;
@@ -161,7 +141,7 @@ int main(const int argc, char **argv)
     std::cout << "  Name: " << tensor_ptr->name << std::endl;
   }
 
-  interpreter->SetNumThreads(thread.value());
+  interpreter->SetNumThreads(thread_count);
 
   std::shared_ptr<nnfw::support::tflite::Session> sess;