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（） - ＆gt; set_visible_device_list（＆＃34; 0＆＃34;）工作，但我不确定。

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

与此类似（针对内存使用限制）： how to limit GPU usage in tensorflow (r1.1) with C++ API

为tensorflow c ++ api的会话选择特定的gpu

1 个答案: