美文网首页工作生活
tensorflow设置GPU

tensorflow设置GPU

作者: ryan_ren | 来源:发表于2019-07-03 13:36 被阅读0次
#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);

  1. 最终成功版
SessionOptions options;
options.config.mutable_gpu_options()->set_visible_device_list("1");
  1. 限制GPU的数目
SessionOptions options;
ConfigProto* config = &options.config;
(*config->mutable_device_count())["GPU"] = 1;

相关文章

网友评论

    本文标题:tensorflow设置GPU

    本文链接:https://www.haomeiwen.com/subject/ysyqhctx.html