TensorFlow使用GPU计算加速

使用方法:tf.device(‘/cpu:0’)或tf.device(‘/gpu:0’)。

实例:

import tensorflow as tf

with tf.device('/cpu:0'):
    a = tf.constant([1.,2.,3.],shape=[3],name='a')
    b = tf.constant([2.,3.,4.],shape=[3],name='b')
with tf.device('/gpu:0'):
    c = a + b

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
sess.run(c)

会打印出:

2018-07-10 23:48:12.430066: I tensorflow/core/common_runtime/placer.cc:886] add: (Add)/job:localhost/replica:0/task:0/device:GPU:0

2018-07-10 23:48:12.430081: I tensorflow/core/common_runtime/placer.cc:886] b: (Const)/job:localhost/replica:0/task:0/device:GPU:0

2018-07-10 23:48:12.430087: I tensorflow/core/common_runtime/placer.cc:886] a: (Const)/job:localhost/replica:0/task:0/device:GPU:0

在默认情况下,即使机器有多个cpu,Tensorflow也不会去区分它们,统一使用/cpu:0。

而同一台机器上不同GPU的名称是不同的,如/gpu:0,/gpu:1等。

默认情况下,Tensorflow优先使用GPU。

需要注意的是,在Tensorflow上,不是所有的操作都可以放在GPU上的,如:

import tensorflow as tf

a_cpu = tf.Variable(0,name='a_cpu')
with tf.device('/gpu:0'):
    a_gpu = tf.Variable(0,name='a_gpu')

with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    sess.run(tf.initialize_all_variables())

则会报错:

InvalidArgumentError (see above for traceback): Cannot assign a device for operation 'a_gpu': Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
Colocation Debug Info:
Colocation group had the following types and devices: 
Assign: CPU 
VariableV2: CPU 
Identity: CPU 

Colocation members and user-requested devices:
  a_gpu (VariableV2) /device:GPU:0
  a_gpu/read (Identity) /device:GPU:0
  a_gpu/Assign (Assign) /device:GPU:0

Registered kernels:
  device='CPU'
  device='GPU'; dtype in [DT_INT64]
  device='GPU'; dtype in [DT_DOUBLE]
  device='GPU'; dtype in [DT_FLOAT]
  device='GPU'; dtype in [DT_HALF]

     [[Node: a_gpu = VariableV2[container="", dtype=DT_INT32, shape=[], shared_name="", _device="/device:GPU:0"]()]]

为了避免这个问题,可以在生成Session时指定allow_soft_placement=True,当运算无法在GPU上执行时,会自动将运算放到CPU上。

用法:

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

Tensorflow会默认占用设备所有GPU以及每个GPU上的显存,如果只使用部分GPU可以:

(注:虽然占用所有GPU,但是会优先使用/GPU:0)

#命令行用法
CUDA_VISIBLE_DEVICES=0,1 python demo.py

或者

#在代码中使用
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

TensorFlow默认一次性占用GPU的所有显存,但是也支持动态分配GPU的显存,使得不会一开始就占满所有显存。

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
#也可以直接按固定的比例分配
#config.gpu_options.per_process_gpu_memory_fraction = 0.4
sess = tf.Session(config=config)

总结几个参数

log_device_placement:将运行每一个操作的设备输出到屏幕上。

allow_soft_placement:将GPU上不能运行的运算自动放到CPU上运行。

allow_growth:动态分配GPU显存。

per_process_gpu_memory_fraction:按比例分配GPU显存。

猜你喜欢

转载自blog.csdn.net/a13602955218/article/details/80994126