TensorFlow多GPU并行计算

1. TensorFlow指定特定GPU或者CPU进行计算：

说明：示例计算机为单CPU（编号为0），单GPU（编号为0），安装的TensorFlow为GPU版。

本文的结构如下：

• 默认为GPU #0
• 指定CPU #0
• 指定GPU #1
• 指定GPU #0 + CPU #0

1.1 默认为GPU #0

I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcuda.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcurand.so locally
           

In []: with tf.Session() as sess:
   ...:     matrix1=tf.constant([[,]])
   ...:     matrix2=tf.constant([[],[]])
   ...:     product=tf.matmul(matrix1,matrix2)
   ...:     result=sess.run(product)
   ...:     print result
   ...:
           

I tensorflow/core/common_runtime/gpu/gpu_init.cc:] Found device  with properties: 
name: GeForce GTX 
major:  minor:  memoryClockRate (GHz) 
pciBusID ::
Total memory: GiB
Free memory: GiB
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] DMA:  
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] :   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:] Creating TensorFlow device (/gpu:) -> (device: , name: GeForce GTX , pci bus id: ::)
[[ 12.]]
           

1.2 指定GPU #0

In []: import tensorflow as tf
           

I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcuda.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcurand.so locally
           

In []: with tf.Session() as sess:
   ...:     with tf.device("/gpu:0"):
   ...:         matrix1=tf.constant([[,]])
   ...:         matrix2=tf.constant([[],[]])
   ...:         product=tf.matmul(matrix1,matrix2)
   ...:     result=sess.run(product)
   ...:     print result
   ...:
           

I tensorflow/core/common_runtime/gpu/gpu_init.cc:] Found device  with properties: 
name: GeForce GTX 
major:  minor:  memoryClockRate (GHz) 
pciBusID ::
Total memory: GiB
Free memory: GiB
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] DMA:  
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] :   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:] Creating TensorFlow device (/gpu:) -> (device: , name: GeForce GTX , pci bus id: ::)
[[ 12.]]
           

1.3 指定CPU #0

I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcuda.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcurand.so locally
           

In []: with tf.Session() as sess:
   ...:     with tf.device("/cpu:0"):
   ...:         matrix1=tf.constant([[,]])
   ...:         matrix2=tf.constant([[],[]])
   ...:         product=tf.matmul(matrix1,matrix2)
   ...:     result=sess.run(product)
   ...:     print result
   ...:
           

I tensorflow/core/common_runtime/gpu/gpu_init.cc:] Found device  with properties: 
name: GeForce GTX 
major:  minor:  memoryClockRate (GHz) 
pciBusID ::
Total memory: GiB
Free memory: GiB
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] DMA:  
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] :   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:] Creating TensorFlow device (/gpu:) -> (device: , name: GeForce GTX , pci bus id: ::)
[[ 12.]]
           

1.4 指定GPU #1

In []: with tf.Session() as sess:
   ...:     with tf.device("/gpu:1"):
   ...:         matrix1=tf.constant([[,]])
   ...:         matrix2=tf.constant([[],[]])
   ...:         product=tf.matmul(matrix1,matrix2)
   ...:     result=sess.run(product)
   ...:     print result
   ...:
           

I tensorflow/core/common_runtime/gpu/gpu_device.cc:] Creating TensorFlow device (/gpu:) -> (device: , name: GeForce GTX , pci bus id: ::)
           

InvalidArgumentError     Traceback (most recent call last)

<ipython-input--ab0827> in <module>()
               matrix2=tf.constant([[2.],[2.]])
               product=tf.matmul(matrix1,matrix2)
----> 6     result=sess.run(product)
           print result
      
           

InvalidArgumentError: **Cannot assign a device to node** 'MatMul_2': 
Could not satisfy explicit device specification '/device:GPU:1' 
because no devices matching that specification are registered in this process; 
available devices: /job:localhost/replica:/task:/cpu:, /job:localhost/replica:/task:/gpu:
[[Node: MatMul_2 = MatMul[T=DT_FLOAT, transpose_a=false, transpose_b=false, _device="/device:GPU:1"](Const_4, Const_5)]]
           

说明：因为本机只有一块GPU，编号为0,而我指定该计算在GPU:1中进行，才报错。

1.4 指定GPU #0 + GPU #0

I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcuda.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcurand.so locally
           

In []: with tf.Session() as sess:
   ...:     with tf.device("/cpu:0"):
   ...:         matrix1=tf.constant([[,]])
   ...:         matrix2=tf.constant([[],[]])
   ...:     with tf.device("/gpu:0"):
   ...:         product=tf.matmul(matrix1,matrix2)
   ...:     result=sess.run(product)
   ...:     print result
   ...:
           

I tensorflow/core/common_runtime/gpu/gpu_init.cc:] Found device  with properties: 
name: GeForce GTX 
major:  minor:  memoryClockRate (GHz) 
pciBusID ::
Total memory: GiB
Free memory: GiB
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] DMA:  
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] :   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:] Creating TensorFlow device (/gpu:) -> (device: , name: GeForce GTX , pci bus id: ::)
[[ 12.]]
           

当然，也可以直接这么写：

I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcuda.so locally
I tensorflow/stream_executor/dso_loader.cc:] successfully opened CUDA library libcurand.so locally
           

In []: with tf.device("/cpu:0"):
   ...:     matrix1=tf.constant([[,]])
   ...:     matrix2=tf.constant([[],[]])
   ...:

In []: with tf.device("/gpu:0"):
   ...:     product=tf.matmul(matrix1,matrix2)
   ...:

In []: sess=tf.Session()
           

I tensorflow/core/common_runtime/gpu/gpu_init.cc:] Found device  with properties: 
name: GeForce GTX 
major:  minor:  memoryClockRate (GHz) 
pciBusID ::
Total memory: GiB
Free memory: GiB
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] DMA:  
I tensorflow/core/common_runtime/gpu/gpu_init.cc:] :   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:] Creating TensorFlow device (/gpu:) -> (device: , name: GeForce GTX , pci bus id: ::)
           

In []: result=sess.run(product)

In []: print result
[[ 12.]]
           

In []: sess.close()
           

注意：将节点operation放到 with tf.device(..): 里面,而启动语句或者不需要计算资源的语句放到with的外面