[DeepRec] Add flag for device_placement_optimization. (#21)

Signed-off-by: Tao Peng <[email protected]>

Author: shanshanpt

tensorflow_serving/model_servers/main.cc

@@ -227,6 +227,9 @@ int main(int argc, char** argv) {
                        "oss_access_id"),
       tensorflow::Flag("oss_access_key", &options.oss_access_key,
                        "oss_access_key"),
+      tensorflow::Flag("enable_device_placement_optimization",
+                       &options.enable_device_placement_optimization,
+                       "Use device placement optimization."),
       tensorflow::Flag("use_multi_stream", &options.use_multi_stream,
                         "Use multi-stream or not in session_group")};
 

tensorflow_serving/model_servers/server.cc

@@ -233,6 +233,13 @@ Status CreatePlatformConfigMap(const Server::Options& server_options,
 
   ParseTimelineConfig(server_options);
 
+  if (server_options.enable_device_placement_optimization) {
+    session_bundle_config.mutable_session_config()
+        ->mutable_graph_options()
+        ->mutable_optimizer_options()
+        ->set_device_placement_optimization(true);
+  }
+
   session_bundle_config.mutable_session_config()
       ->mutable_gpu_options()
       ->set_per_process_gpu_memory_fraction(
@@ -294,6 +301,13 @@ Status CreatePlatformConfigMapV2(const Server::Options& server_options,
 
   ParseTimelineConfig(server_options);
 
+  if (server_options.enable_device_placement_optimization) {
+    model_session_config->mutable_session_config()
+        ->mutable_graph_options()
+        ->mutable_optimizer_options()
+        ->set_device_placement_optimization(true);
+  }
+
   // session num
   model_session_config->set_session_num(
       server_options.session_num_per_group);

tensorflow_serving/model_servers/server.h

@@ -100,6 +100,8 @@ class Server {
     tensorflow::string oss_access_id = "";
     tensorflow::string oss_access_key = "";
 
+    bool enable_device_placement_optimization = false;
+
     Options();
   };