为TensorFlow C ++ API的会话选择特定的GPU

Question

How could I ask tensorflow use specific gpu to do the inference?

Part of the source codes

std::unique_ptr<tensorflow::Session> session;  
Status const load_graph_status = LoadGraph(graph_path, &session);
if (!load_graph_status.ok()) {
   LOG(ERROR) << "LoadGraph ERROR!!!!"<< load_graph_status;
   return -1;
}

std::vector<Tensor> resized_tensors;
Status const read_tensor_status = ReadTensorFromImageFile(image_path, &resized_tensors);
if (!read_tensor_status.ok()) { 
    LOG(ERROR) << read_tensor_status;
    return -1;
}

std::vector<Tensor> outputs;
Status run_status = session->Run({{input_layer, resized_tensor}},
                                   output_layer, {}, &outputs);

So far so good, but tensorflow always select the same gpu when I execute Run, do I have a way to specify which gpu to execute?

In case you need complete source codes, I placed them at pastebin

Edit : Looks like options.config.mutable_gpu_options()->set_visible_device_list("0") work, but I am not sure.

Answer 1

Turns out in the C++ API there are a series of (nested) structs: tensorflow::SessionOptions , tensorflow::ConfigProto , and tensorflow::GPUOptions . The latter contains a method called set_visible_device_list(::std::string&& value) which you can select the GPU you would like:

  auto options = tensorflow::SessionOptions();
  options.config.mutable_gpu_options()->set_visible_device_list("0");
  // session_ is a unique_ptr to a tensorflow::Session
  session_->reset(tensorflow::NewSession(options));

Similar to this (for memory usage restriction): how to limit GPU usage in tensorflow (r1.1) with C++ API