tensorflow在指定设备时,注意:
-
CPU在内部计算的时候,不会涉及到它是第几个CPU,就是"/cpu:0"表示使用机器CPU运算
-
而GPU会。比如:
"/gpu:0"表示使用第一块GPU运算
"/gpu:1"表示使用第二块GPU运算
"/gpu:2"表示使用第三块GPU运算 -
Device mapping只会显示GPU的信息,如果是CPU就没有了。
tf.device(设备)
import tensorflow as tf
# tensorflow指定运算设备
with tf.device('/cpu:0'):
# 这个代码块中定义的操作,会在tf.device给定的设备上运行
# 有一些操作,是不会在GPU上运行的(注意了) 所以建议如果安装的是GPU版本,把参数 allow_soft_placement设为True
# 如果安装的tensorflow是CPU版本,是没法指定运行环境的,因为指不指定都在仅有的CPU上面运算
a = tf.Variable([1, 2, 3], dtype=tf.int32, name='a')
b = tf.Variable(2, dtype=tf.int32, name='b')
c = tf.add(a, b, name='ab')
with tf.device('/gpu:0'): # 我没有GPU环境,运算没找到GPU,就在CPU上运算了
d = tf.Variable([2, 8, 13], dtype=tf.int32, name='d')
e = tf.constant(2, dtype=tf.int32, name='e')
f = d+e
g = c+f
with tf.Session(config=tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)) as sess:
tf.global_variables_initializer().run()
print(sess.run(g))