[英]Select specific gpu for the session of tensorflow c++ api
我怎样才能让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”)工作,但我不确定。
事实证明,在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使用
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.