TensorFlow入门(五、指定GPU运算)

一般情况下,下载的TensorFlow版本如果是GPU版本,在运行过程中TensorFlow能自动检测。如果检测到GPU,TensorFlow会默认利用找到的第一个GPU来执行操作。如果机器上有超过一个可用的GPU,除第一个之外的其他GPU默认是不参与计算的。如果想让TensorFlow使用这些GPU执行操作,需要将运算OP明确指派给它们执行。

指派特定的CPU或GPU执行操作的方法有两种:

①使用with……device语句来指派,示例代码如下:

#使用with语法建立session
with tf.Session() as sess:
    #指定CPU
    with tf.device('/cpu:0'):
        #计算具体数值  使用feed机制
        print("相加: %i" % sess.run(add,feed_dict = {a:3,b:4}))
        print("相乘: %i" % sess.run(mul,feed_dict = {a:3,b:4}))
        
        #使用fetch机制
        print(sess.run([mul,add],feed_dict = {a:3,b:4}))

从代码中可以看到,设备用字符串进行标识,目前支持的设备包括以下几种:

                        ①cpu:0        机器的cpu

                        ②gpu:0        机器的第一个gpu,如果有的话

                        ③gpu:1        机器的第二个gpu,以此类推

②通过tf.ConfigProto来构建一个config,在config中指定相关的GPU,并且在session中传入参数config="自己创建的config"来指定GPU操作

        tf.ConfigProto函数的参数具体如下:

#log_device_placement = True    是否打印设备分配日志
#allow_soft_placement = True    如果指定的设备不存在,允许TF自动分配设备


myconfig = tf.ConfigProto(log_device_placemnet = True,allow_soft_placement = True)
session = tf.Session(config = myconfig)

        通过tf.ConfigProto函数生成config后,还可以设置其属性来分配GPU的运算资源,如下代码的用图就是按需分配:

myconfig = tf.ConfigProto(log_device_placement = True, allow_soft_placement = true)
config.gpu_options.allow_growth = True
session = tf.Session(config = myconfig)

        使用allow_growth刚开始会分配少量的GPU容量,然后按需慢慢地增加,

由于不会释放内存,会导致碎片。

        也可以放在config创建的时候指定,如:

gpu_options = tf.GPUOptions(allow_growth = True)
myconfig = tf.ConfigProto(gpu_options = gpu_options)

        还可以给GPU分配固定大小的计算资源,如:

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = 0.5)

        含义为分配给TensorFlow的GPU显存大小为GPU实际显存*0.5

完整示例代码如下:

import tensorflow.compat.v1 as tf

tf.compat.v1.disable_eager_execution()

#通过占位符定义变量
a = tf.compat.v1.placeholder(tf.int16)
b = tf.compat.v1.placeholder(tf.int16)

#a与b相加
add = tf.add(a,b)
#a与b相乘
mul = tf.multiply(a,b)

gpu_options = tf.GPUOptions(allow_growth = True,per_process_gpu_memory_fraction = 0.5)
myconfig = tf.ConfigProto(gpu_options = gpu_options)

#使用with语法建立session
with tf.Session(config = myconfig) as sess:
    #计算具体数值  使用feed机制
    print("相加: %i" % sess.run(add,feed_dict = {a:3,b:4}))
    print("相乘: %i" % sess.run(mul,feed_dict = {a:3,b:4}))
    
    #使用fetch机制
    print(sess.run([mul,add],feed_dict = {a:3,b:4}))

运行结果如下: