在构建计算图的时候,规定好在哪个设备上跑
#create a graph.
with tf.device('/gpu:2'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name = 'a' )
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name = 'b' )
c = tf.matmul(a, b)
#creates a session with log_device_placement set to True.
sess = tf.Session(config = tf.ConfigProto(log_device_placement = True))
#Run the op.
print sess.run(c)
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////