为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使用