import numpy
import tensorflow as tf
with tf.device('/gpu:0'):
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)
加入config=tf.ConfigProto(log_device_placement=True)这个参数后在执行的过程中就能打印具体是在CPU上执行还是GPU上执行了。
hank@hank-desktop:~/Study/CV$ python ./test_tensorflow_run_with_gpu_cpu.py
2018-08-26 10:25:57.854091: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:964] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2018-08-26 10:25:57.854738: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1406] Found device 0 with properties:
name: GeForce GTX 750 Ti major: 5 minor: 0 memoryClockRate(GHz): 1.0845
pciBusID: 0000:01:00.0
totalMemory: 3.95GiB freeMemory: 3.91GiB
2018-08-26 10:25:57.854831: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1485] Adding visible gpu devices: 0
2018-08-26 10:25:58.200876: I tensorflow/core/common_runtime/gpu/gpu_device.cc:966] Device interconnect StreamExecutor with strength 1 edge matrix:
2018-08-26 10:25:58.200957: I tensorflow/core/common_runtime/gpu/gpu_device.cc:972] 0
2018-08-26 10:25:58.200980: I tensorflow/core/common_runtime/gpu/gpu_device.cc:985] 0: N
2018-08-26 10:25:58.201344: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1098] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 3640 MB memory) -> physical GPU (device: 0, name: GeForce GTX 750 Ti, pci bus id: 0000:01:00.0, compute capability: 5.0)
Device mapping:
/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce GTX 750 Ti, pci bus id: 0000:01:00.0, compute capability: 5.0
2018-08-26 10:25:58.236122: I tensorflow/core/common_runtime/direct_session.cc:291] Device mapping:
/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce GTX 750 Ti, pci bus id: 0000:01:00.0, compute capability: 5.0
MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0
2018-08-26 10:25:58.237332: I tensorflow/core/common_runtime/placer.cc:923] MatMul: (MatMul)/job:localhost/replica:0/task:0/device:GPU:0
a: (Const): /job:localhost/replica:0/task:0/device:GPU:0
2018-08-26 10:25:58.237377: I tensorflow/core/common_runtime/placer.cc:923] a: (Const)/job:localhost/replica:0/task:0/device:GPU:0
b: (Const): /job:localhost/replica:0/task:0/device:GPU:0
2018-08-26 10:25:58.237454: I tensorflow/core/common_runtime/placer.cc:923] b: (Const)/job:localhost/replica:0/task:0/device:GPU:0
[[22. 28.]
[49. 64.]]
