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.
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
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.