相关疑难解决方法(0)

我在一个共享计算资源的环境中工作,也就是说,我们有一些服务器机器配备了几个Nvidia Titan X GPU.

对于小到中等大小的型号,12GB的Titan X通常足以让2-3人在同一GPU上同时进行训练.如果模型足够小以至于单个模型没有充分利用Titan X的所有计算单元,那么与在另一个训练过程之后运行一个训练过程相比,这实际上可以导致加速.即使在并发访问GPU确实减慢了单个培训时间的情况下,仍然可以灵活地让多个用户同时在GPU上运行.

TensorFlow的问题在于,默认情况下,它在启动时会在GPU上分配全部可用内存.即使对于一个小的2层神经网络,我也看到12 GB的Titan X已用完.

有没有办法让TensorFlow只分配4GB的GPU内存,如果有人知道这个数量对于给定的模型来说足够了？

我怎样才能让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”）工作，但我不确定。