☰
Navigation :

GPU

TensorFlow 支持 CPU 和 GPU 这两种设备，标识设备的方法为:

记录设备指派情况

通过设置 log_device_placement 选项来记录 operations 和 Tensor 被指派到哪个设备上运行：

import tensorflow as tf

a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
id: 0000:05:00.0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
MatMul: /job:localhost/replica:0/task:0/gpu:0
[[22.  28.]
 [49.  64.]]

手动指派设备

import tensorflow as tf

with tf.device('/gpu:2'):
 a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
 b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True))
print sess.run(c)

为了避免出现指定设备不存在的情况, 可以在创建的 session 里把参数 allow_soft_placement 设置为 True, 这样 tensorFlow 会自动选择一个存在并且支持的设备来运行 operation.

示例

基于多个设备计算圆周率π的方法

import tensorflow as tf
import numpy as np
c = []

#Distribute the work between the GPUs
for d in ['/gpu:0', '/gpu:1', '/gpu:2', '/gpu:3']:
    #Generate the random 2D samples
    i=tf.constant(np.random.uniform(size=10000), shape=[5000,2])
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        #Calculate the euclidean distance to the origin
        distances=tf.reduce_sum(tf.pow(i,2),1)
        #Sum the samples inside the circle
        tempsum = sess.run(tf.reduce_sum(tf.cast(tf.greater_equal(tf.cast(1.0,tf.float64),distances),tf.float64)))
        #append the current result to the results array
        c.append(tempsum)
    #Do the final ratio calculation on the CPU
    with tf.device('/cpu:0'):
        with tf.Session() as sess:
            sum = tf.add_n(c)
            print (sess.run(sum/20000.0)*4.0)