File make-infer-device-configureable.patch of Package cloud-native-demos

Overview Repositories Revisions Requests Users Attributes Meta

File make-infer-device-configureable.patch of Package cloud-native-demos

diff --git a/elastic_inference/apps/clcn/nn/inferengine.py b/elastic_inference/apps/clcn/nn/inferengine.py
index 69e0ff4..b4701af 100644
--- a/elastic_inference/apps/clcn/nn/inferengine.py
+++ b/elastic_inference/apps/clcn/nn/inferengine.py
@@ -1,6 +1,7 @@
 """
 Openvino based inference engine
 """
+import os
 import logging
 import time
 import cv2
@@ -129,11 +130,34 @@ class OpenVinoInferEngineTask(InferEngineTask):
             report_metric_fn)
         LOG.info("Model dir: %s", model_dir)
         LOG.info("Model name: %s", model_name)
-        self._plugin = self._init_openvino_cpu_plugin()
+        plugin_device = self.get_env("INFER_DEVICE", "CPU")
+        if plugin_device.lower() == "cpu":
+            self._plugin = self._init_openvino_cpu_plugin()
+        elif plugin_device.lower() == "myriad":
+            self._plugin = self._init_openvino_myriad_plugin()
+        else:
+            LOG.error("Unknown device [%s], fallback to CPU", plugin_device)
+            self._plugin = self._init_openvino_cpu_plugin()
         self._nn = NNFactory.get_detection(model_dir, model_name)
         self._nn.load()
         self._exec = self._plugin.load(network=self._nn.net)
 
+    @staticmethod
+    def get_env(key, default=None):
+        """
+        Get environment variable.
+        """
+        if key not in os.environ:
+            LOG.warning("Could not find the key %s in environment, use "
+                        "default value %s", key, str(default))
+            return default
+        return os.environ[key]
+
+    @staticmethod
+    def _init_openvino_myriad_plugin():
+        plugin = IEPlugin(device="MYRIAD")
+        return plugin
+
     @staticmethod
     def _init_openvino_cpu_plugin():
         plugin = IEPlugin(device="CPU")

Places

File make-infer-device-configureable.patch of Package cloud-native-demos

Places