TensorFlow 支持 CPU 和 GPU。它也支持分布式计算。可以在一个或多个计算机系统的多个设备上使用 TensorFlow。
TensorFlow 将支持的 CPU 设备命名为“/device:CPU:0”(或“/cpu:0”),第 i 个 GPU 设备命名为“/device:GPU:I”(或“/gpu:I”)。
具体做法
import tensorflow as tf
手动选择 CPU 进行操作:
import tensorflow as tf
with tf.device(’/cpu:0’):
v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name=‘v1’)
v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name=‘v2’)
sumV12 = v1 + v2
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
print (sess.run(sumV12))
手动选择一个GPU 来操作:
使用 tf.device(’/gpu:1’) 指定Session在第二块GPU上运行:
import tensorflow as tf
with tf.device(’/gpu:1’):
v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name=‘v1’)
v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name=‘v2’)
sumV12 = v1 + v2
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
print sess.run(sumV12)
ConfigProto() 中参数 log_device_placement=True 会打印出执行操作所用的设备
如果安装的是GPU版本的tensorflow,机器上有支持的GPU,也正确安装了显卡驱动、CUDA和cuDNN,默认情况下,Session会在GPU上运行,可以通过改上面那个参数
手动选择多个GPU:
拓展
函数 tf.device() 选择设备(CPU 或 GPU)。with 块确保设备被选择并用于其操作。with 块中定义的所有变量、常量和操作将使用在 tf.device() 中选择的设备。
会话配置使用 tf.ConfigProto 进行控制。通过设置 allow_soft_placement 和 log_device_placement 标志,告诉 TensorFlow 在指定的设备不可用时自动选择可用的设备,并在执行会话时给出日志消息作为描述设备分配的输出。