為TensorFlow C ++ API的會話選擇特定的GPU

Question

我怎樣才能讓tensorflow使用特定的gpu進行推斷？

部分源代碼

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);

到目前為止一切都很好，但是當我執行Run時，tensorflow總是選擇相同的gpu，我是否有辦法指定要執行的gpu？

如果您需要完整的源代碼，我將它們放在pastebin上

編輯：看起來options.config.mutable_gpu_options（）-> set_visible_device_list（“ 0”）工作，但我不確定。

Answer 1

事實證明，在C ++ API中，存在一系列（嵌套的）結構： tensorflow::SessionOptions ， tensorflow::ConfigProto和tensorflow::GPUOptions 。 后者包含一個名為set_visible_device_list(::std::string&& value) ，您可以選擇所需的GPU：

  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));

與此類似（用於內存使用限制）： 如何使用C ++ API在Tensorflow（r1.1）中限制GPU使用