TensorFlow指定GPU设备、分配显存

版权声明:本文为博主原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。
本文链接:https://blog.csdn.net/Im_Chenxi/article/details/80252033

一、使用GPU加速运算

        1.TensorFlow程序中可以通过tf.device函数来指定运行的设备。

          CPU名称为"/cpu:0",即使有多个CPU,所有的CPU都使用/cpu:0作为名称。

          GPU则不同,是"/gpu:n",n可替换。n=0,即是第一个GPU,以此类推。

        2.观察每个运算所用的设备:

          使用log_device_placement传入会话当中的配置参数。具体代码:

            

import tensorflow as tf
sess=tf.Session(config=tf.ConfigProto(log_devic_placement=True))
sess.run(...)    

          运行程序后会显示每一个操作所用的设备显示到屏幕上:

          

          如果没有明确指定运行设备,那么TensorFlow会优先选择GPU。

        2.指定运算到指定的设备

          需要tf.device来手工指定。代码如下:

          

          得到如下结果:

          

          注意:有些操作是不能在GPU上完成,否则会报错。

              比如:tf.Variable()操作只支持实数型数据的参数在GPU上运行,整数型就会报错。

              具体哪些不能在GPU上运行,参考:官方部分文档

        为了避免不知道某个操作是否确定在GPU上运行,可指定GPU设备后,在构建会话时,传入参数:                allow_soft_placement=True,使得不能在GPU设备上完成的自动移到CPU上完成,具体代码:

              

sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

二、TensorFlow动态分配GPU和显存

          TensorFlow会默认占用所有GPU以及每个GPU的所有显存。

          1.如果只想使用部分GPU,2种方法进行设置:

            第一种:

            运行.py程序时指定:CUDA_VISIBLE_DEVICES=0, 1 python xx.py

            这时候就会只有第一块、第二块GPU会使用。

            第二种:

            在程序中指定:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
#只使用第三块GPU

          2.动态分配显存

          让TensorFlow按需分配显存:      

config = tf.ConfigProto()
config.gpu_options.allow_growth = True

          直接指定显存比例:

config.gpu_options.per_process_gpu_memory_fraction = 0.4
#占用GPU40%

           再传入会话当中:

session = tf.Session(config=config)

展开阅读全文

没有更多推荐了,返回首页