-
#include <stdlib.h>
setenv("CUDA_VISIBLE_DEVICES", "", 1);
-
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/graph/default_device.h"
Session* session;
GraphDef graph_def;
SessionOptions opts;
graph::SetDefaultDevice("/gpu:0", &graph_def);
# 限制每个GPU的占用率
opts.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.5);
# 设置显卡内存随实际情况而增长
opts.config.mutable_gpu_options()->set_allow_growth(true);
- 最终成功版
SessionOptions options;
options.config.mutable_gpu_options()->set_visible_device_list("1");
- 限制GPU的数目
SessionOptions options;
ConfigProto* config = &options.config;
(*config->mutable_device_count())["GPU"] = 1;
