C++编译的Tensorflow进行显存限制

tensorflow进行模型时会将显存占满,导致不能与其他模型并行。

在Python中的Tensorflow显存设置代码

config=tf.ConfigProto();
config.gpu_options.per_process_gpu_memory_fraction=0.5
set_session(tf.Session(config=config))

在C++中

GPUOption类中有一个per_process_gpu_memory_fraction对象。其默认值为1。所以会每次将显存占满,将改值修改即可。

    Session* session;

    ConfigProto configProto;
    //新建一个gpuoption对象
    GPUOptions gpuOptions;
    //设置显存
    gpuOptions.set_per_process_gpu_memory_fraction(0.3);
    //修改configProto中的gpuoption对象
    configProto.set_allocated_gpu_options(&gpuOptions);
    
    //新建一个sessionOption对象并将configProto对象赋值给它
    SessionOptions sessionOptions;
    sessionOptions.config=configProto;
    Status status = NewSession(sessionOptions, &session);//创建新会话Session

SessionOption类中含有ConfigProto的对象,若不需要对显存进行修改,则直接调用SessionOptions的默认构造函数即可。

Status status = NewSession(SessionOptions(), &session);

一开始是想使用configProto.gpuoption()函数来获得gpuoption对象,并对它进行修改。但是发现该函数返回的是一个const对象。。。。

后来使用ConfigProto.set_allocated_gpu_options()函数完成。

不加显存限制结果如下:

加显存限制结果:

扫描二维码关注公众号,回复: 8519178 查看本文章

因为出来Tensorflow以外还有其他模型,所以显卡数值不是完全的0.3。但是可以看出明显限制显存了。可以同时加载多个模型了。

发布了58 篇原创文章 · 获赞 64 · 访问量 1万+

猜你喜欢

转载自blog.csdn.net/ThorKing01/article/details/103068085