TensorFlow程式可以通過tf.device函式來指定執行每一個操作的裝置，這個裝置可以是本地的CPU或者GPU，也可以是某一臺遠端的伺服器。但在本文中只關心本地的裝置。TensorFlow會給每一個可用的裝置一個名稱，tf.device函式可以通過裝置的名稱來指定執行運算的裝置。比如CPU在TensorFlow中的名稱為/cpu:0。在預設情況下，即使機器有多個CPU，TensorFlow也不會區分它們，所有的CPU都使用/cpu:0作為名稱。而一臺機器上不同GPU的名稱是不同的，第n個GPU在TensorFlow中的名稱為/gpu:n。比如第一個GPU的名稱為/gpu:0，第二個GPU名稱為/gpu:1，以此類推。

TensorFlow提供了一個快捷的方式來檢視執行每一個運算的裝置。在生成會話時，可以通過設定log_device_placement引數來列印執行每一個運算的裝置。以下程式展示瞭如何使用log_device_placement這個引數。

import tensorflow as tf

a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
c = a + b
# 通過log_device_placement引數來輸出執行每一個運算的裝置。
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
在沒有GPU的機器上執行以上程式碼可以得到類似以下的輸出：
Device mapping: no known devices.

add: (Add): /job:localhost/replica:0/task:0/cpu:0
b: (Const): /job:localhost/replica:0/task:0/cpu:0
a: (Const): /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''
 

在以上程式碼中，TensorFlow程式生成會話時加入了引數log_device_placement=True，所以程式會將執行每一個操作的裝置輸出到螢幕。於是除了可以看到最後的計算結果，還可以看到類似“add: /job:localhost/replica:0/task:0/cpu:0”這樣的輸出。這些輸出顯示了執行每一個運算的裝置。比如加法操作add是通過CPU來執行的，因為它的裝置名稱中包含了/cpu:0。

在配置好GPU環境的TensorFlow中，如果操作沒有明確地指定執行裝置，那麼TensorFlow會優先選擇GPU。比如將以上程式碼在亞馬遜（Amazon Web Services, AWS）的 g2.8xlarge例項上執行時，會得到類似以下的執行結果。

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: (Add): /job:localhost/replica:0/task:0/gpu:0
b: (Const): /job:localhost/replica:0/task:0/gpu:0
a: (Const): /job:localhost/replica:0/task:0/gpu:0
[ 2.  4.  6.]

從以上輸出可以看到在配置好GPU環境的TensorFlow中，TensorFlow會自動優先將運算放置在GPU上。不過，儘管g2.8xlarge例項有4個GPU，在預設情況下，TensorFlow只會將運算優先放到/gpu:0上。於是可以看見在以上程式中，所有的運算都被放在了/gpu:0上。如果需要將某些運算放到不同的GPU或者CPU上，就需要通過tf.device來手工指定。以下程式給出了一個通過tf.device手工指定執行裝置的樣例。

import tensorflow as tf

# 通過tf.device將運算指定到特定的裝置上。
with tf.device('/cpu:0'):
   a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
   b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')

with tf.device('/gpu:1'):
    c = a + b

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
在AWS g2.8xlarge例項上執行上述程式碼可以得到以下結果：
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: (Add): /job:localhost/replica:0/task:0/gpu:1
b: (Const): /job:localhost/replica:0/task:0/cpu:0
a: (Const): /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上程式碼中可以看到生成常量a和b的操作被載入到了CPU上，而加法操作被放到了第二個GPU“/gpu:1”上。在TensorFlow中，不是所有的操作都可以被放在GPU上，如果強行將無法放在GPU上的操作指定到GPU上，那麼程式將會報錯。以下程式碼給出了一個報錯的樣例。

import tensorflow as tf

# 在CPU上執行tf.Variable
a_cpu = tf.Variable(0, name="a_cpu")

with tf.device('/gpu:0'):
    # 將tf.Variable強制放在GPU上。
    a_gpu = tf.Variable(0, name="a_gpu")

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
sess.run(tf.initialize_all_variables())

'''
執行以上程式將會報出以下錯誤：
tensorflow.python.framework.errors.InvalidArgumentError: Cannot assign a device to node 'a_gpu': Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
Colocation Debug Info:
Colocation group had the following types and devices: 
Identity: CPU 
Assign: CPU 
Variable: CPU 
[[Node: a_gpu = Variable[container="", dtype=DT_INT32, shape=[], shared_ name="", _device="/device:GPU:0"]()]]
'''

不同版本的TensorFlow對GPU的支援不一樣，如果程式中全部使用強制指定裝置的方式會降低程式的可移植性。在TensorFlow的kernel中定義了哪些操作可以跑在GPU上。比如可以在variable_ops.cc程式中找到以下定義。

# define REGISTER_GPU_KERNELS(type)                                       \
    REGISTER_KERNEL_BUILDER(                                                 \
      Name("Variable").Device(DEVICE_GPU).TypeConstraint<type>("dtype"),\
      VariableOp);                                                                 \
   …
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS);

在這段定義中可以看到GPU只在部分資料型別上支援tf.Variable操作。如果在TensorFlow程式碼庫中搜索呼叫這段程式碼的巨集TF_CALL_GPU_NUMBER_TYPES，可以發現在GPU上，tf.Variable操作只支援實數型（float16、float32和double）的引數。而在報錯的樣例程式碼中給定的引數是整數型的，所以不支援在GPU上執行。為避免這個問題，TensorFlow在生成會話時可以指定allow_soft_placement引數。當allow_soft_placement引數設定為True時，如果運算無法由GPU執行，那麼TensorFlow會自動將它放到CPU上執行。以下程式碼給出了一個使用allow_soft_placement引數的樣例。

import tensorflow as tf

a_cpu = tf.Variable(0, name="a_cpu")
with tf.device('/gpu:0'):
    a_gpu = tf.Variable(0, name="a_gpu")

# 通過allow_soft_placement引數自動將無法放在GPU上的操作放回CPU上。
sess = tf.Session(config=tf.ConfigProto(
    allow_soft_placement=True, log_device_ placement=True))
sess.run(tf.initialize_all_variables())

'''
執行上面這段程式可以得到以下結果：
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0
a_gpu: /job:localhost/replica:0/task:0/cpu:0
a_gpu/read: /job:localhost/replica:0/task:0/cpu:0
a_gpu/Assign: /job:localhost/replica:0/task:0/cpu:0
init/NoOp_1: /job:localhost/replica:0/task:0/gpu:0
a_cpu: /job:localhost/replica:0/task:0/cpu:0
a_cpu/read: /job:localhost/replica:0/task:0/cpu:0
a_cpu/Assign: /job:localhost/replica:0/task:0/cpu:0
init/NoOp: /job:localhost/replica:0/task:0/gpu:0
init: /job:localhost/replica:0/task:0/gpu:0
a_gpu/initial_value: /job:localhost/replica:0/task:0/gpu:0
a_cpu/initial_value: /job:localhost/replica:0/task:0/cpu:0

從輸出的日誌中可以看到在生成變數a_gpu時，無法放到GPU上的運算被自動調整到了CPU上（比如a_gpu和a_gpu/read），而可以被GPU執行的命令（比如a_gpu/initial_value）依舊由GPU執行。
'''

雖然GPU可以加速TensorFlow的計算，但一般來說不會把所有的操作全部放在GPU上。一個比較好的實踐是將計算密集型的運算放在GPU上，而把其他操作放到CPU上。GPU是機器中相對獨立的資源，將計算放入或者轉出GPU都需要額外的時間。而且GPU需要將計算時用到的資料從記憶體複製到GPU裝置上，這也需要額外的時間。TensorFlow可以自動完成這些操作而不需要使用者特別處理，但為了提高程式執行的速度，使用者也需要儘量將相關的運算放在同一個裝置上。

TensorFlow預設會佔用裝置上的所有GPU以及每個GPU的所有視訊記憶體。如果在一個TensorFlow程式中只需要使用部分GPU，可以通過設定CUDA_VISIBLE_DEVICES環境變數來控制。以下樣例介紹瞭如何在執行時設定這個環境變數。

# 只使用第二塊GPU（GPU編號從0開始）。在demo_code.py中，機器上的第二塊GPU的
# 名稱變成/gpu:0，不過在執行時所有/gpu:0的運算將被放在第二塊GPU上。
CUDA_VISIBLE_DEVICES=1 python demo_code.py
# 只使用第一塊和第二塊GPU。
CUDA_VISIBLE_DEVICES=0,1 python demo_code.py

TensorFlow也支援在程式中設定環境變數，以下程式碼展示瞭如何在程式中設定這些環境變數。

import os

# 只使用第三塊GPU。
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

雖然TensorFlow預設會一次性佔用一個GPU的所有視訊記憶體，但是TensorFlow也支援動態分配GPU的視訊記憶體，使得一塊GPU上可以同時執行多個任務。下面給出了TensorFlow動態分配視訊記憶體的方法。

config = tf.ConfigProto()

# 讓TensorFlow按需分配視訊記憶體。

config.gpu_options.allow_growth = True

# 或者直接按固定的比例分配。以下程式碼會佔用所有可使用GPU的40%視訊記憶體。
# config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)

作者：博文視點
連結：https://www.jianshu.com/p/26ac409dfb38
來源：簡書
簡書著作權歸作者所有，任何形式的轉載都請聯絡作者獲得授權並註明出處。