TensorFlow使用GPU計算加速
阿新 • • 發佈:2018-11-21
使用方法: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視訊記憶體。