本頁面由 Cloud Translation API 翻譯而成。
Switch to English

使用GPU

在TensorFlow.org上查看 在Google Colab中運行 在GitHub上查看源代碼 下載筆記本

TensorFlow代碼和tf.keras模型將透明地在單個GPU上運行,而無需更改代碼。

在一台或多台機器上的多個GPU上運行的最簡單方法是使用Distribution Strategies

本指南適用於嘗試過這些方法並發現他們需要對TensorFlow如何使用GPU進行細粒度控制的用戶。

建立

確保您已安裝最新的TensorFlow gpu版本。

 import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
 
Num GPUs Available:  2

總覽

TensorFlow支持在各種類型的設備上運行計算,包括CPU和GPU。它們用字符串標識符表示,例如:

  • "/device:CPU:0" :計算機的CPU。
  • "/GPU:0" :TensorFlow可以看到的機器上第一個GPU的簡寫形式。
  • "/job:localhost/replica:0/task:0/device:GPU:1" :TensorFlow可以看到的機器第二個GPU的全限定名稱。

如果TensorFlow操作同時具有CPU和GPU實施,則默認情況下,當將操作分配給設備時,GPU設備將被賦予優先級。例如, tf.matmul具有CPU和GPU內核。在具有設備CPU:0GPU:0 ,除非您明確要求在其他設備上運行它,否則將選擇GPU:0設備運行tf.matmul

記錄設備放置

要找出您的操作和張量分配給哪些設備,請將tf.debugging.set_log_device_placement(True)作為程序的第一條語句。啟用設備放置日誌記錄將導致打印任何張量分配或操作。

 tf.debugging.set_log_device_placement(True)

# Create some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)
 
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

上面的代碼將顯示MatMul op在GPU:0上執行的指示。

手動放置設備

如果您希望某個特定的操作在您選擇的設備上運行,而不是自動為您選擇的設備,則可以with tf.device一起使用來創建設備上下文,並且該上下文中的所有操作都將在同一指定設備上運行。

 tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

c = tf.matmul(a, b)
print(c)
 
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

您將看到現在將ab分配給CPU:0 。由於未為MatMul操作明確指定設備,因此TensorFlow運行時將根據操作和可用設備(在此示例中為GPU:0 )選擇一個設備,並在需要時自動在設備之間複製張量。

限制GPU內存的增長

默認情況下,TensorFlow會映射該進程可見的幾乎所有GPU的所有GPU內存(取決於CUDA_VISIBLE_DEVICES )。這樣做是為了通過減少內存碎片來更有效地使用設備上相對寶貴的GPU內存資源。要將TensorFlow限制為一組特定的GPU,我們使用tf.config.experimental.set_visible_devices方法。

 gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
  except RuntimeError as e:
    # Visible devices must be set before GPUs have been initialized
    print(e)
 
2 Physical GPUs, 1 Logical GPU

在某些情況下,希望該過程僅分配可用內存的子集,或者僅增加該過程所需的內存使用量。 TensorFlow提供了兩種方法來控制它。

第一種選擇是通過調用tf.config.experimental.set_memory_growth來打開內存增長,該嘗試嘗試僅分配運行時分配所需的GPU內存:它開始分配很少的內存,並且隨著程序的運行和需要更多的GPU內存,我們擴展了分配給TensorFlow進程的GPU內存區域。請注意,我們不會釋放內存,因為它可能導致內存碎片。要打開特定GPU的內存增長,請在分配任何張量或執行任何操作之前使用以下代碼。

 gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)
 
2 Physical GPUs, 2 Logical GPUs

啟用此選項的另一種方法是將環境變量TF_FORCE_GPU_ALLOW_GROWTH設置為true 。此配置是特定於平台的。

第二種方法是使用tf.config.experimental.set_virtual_device_configuration配置虛擬GPU設備, tf.config.experimental.set_virtual_device_configuration要在GPU上分配的總內存設置硬限制。

 gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
 
2 Physical GPUs, 2 Logical GPUs

如果您想真正綁定TensorFlow進程可用的GPU內存量,這將很有用。當GPU與其他應用程序(例如工作站GUI)共享GPU時,這是本地開發的常見做法。

在多GPU系統上使用單個GPU

如果系統中有多個GPU,則默認情況下將選擇ID最低的GPU。如果要在其他GPU上運行,則需要明確指定首選項:

 tf.debugging.set_log_device_placement(True)

try:
  # Specify an invalid GPU device
  with tf.device('/device:GPU:2'):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    c = tf.matmul(a, b)
except RuntimeError as e:
  print(e)
 
/job:localhost/replica:0/task:0/device:GPU:2 unknown device.

如果您指定的設備不存在,您將得到RuntimeError.../device:GPU:2 unknown device

如果您希望TensorFlow在沒有指定的設備不存在的情況下自動選擇一個現有的受支持的設備來運行操作,則可以調用tf.config.set_soft_device_placement(True)

 tf.config.set_soft_device_placement(True)
tf.debugging.set_log_device_placement(True)

# Creates some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)
 
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

使用多個GPU

針對多個GPU的開發將使模型可以使用其他資源進行擴展。如果在具有單個GPU的系統上進行開發,我們可以使用虛擬設備模擬多個GPU。這樣可以輕鬆測試多GPU設置,而無需其他資源。

 gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Create 2 virtual GPUs with 1GB memory each
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
         tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
 
2 Physical GPU, 3 Logical GPUs

一旦我們有多個邏輯GPU可供運行時使用,就可以通過tf.distribute.Strategy或手動放置來使用多個GPU。

使用tf.distribute.Strategy

使用多個GPU的最佳實踐是使用tf.distribute.Strategy 。這是一個簡單的示例:

 tf.debugging.set_log_device_placement(True)

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
  inputs = tf.keras.layers.Input(shape=(1,))
  predictions = tf.keras.layers.Dense(1)(inputs)
  model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
  model.compile(loss='mse',
                optimizer=tf.keras.optimizers.SGD(learning_rate=0.2))
 
Executing op RandomUniform in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Sub in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Add in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op Reshape in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1

該程序將在每個GPU上運行模型的副本,在它們之間拆分輸入數據,也稱為“ 數據並行性 ”。

有關分發策略的更多信息,請在此處查看指南。

手動放置

tf.distribute.Strategy通過在設備之間複製計算來在tf.distribute.Strategy運行。您可以通過在每個GPU上構建模型來手動實現複製。例如:

 tf.debugging.set_log_device_placement(True)

gpus = tf.config.experimental.list_logical_devices('GPU')
if gpus:
  # Replicate your computation on multiple GPUs
  c = []
  for gpu in gpus:
    with tf.device(gpu.name):
      a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
      c.append(tf.matmul(a, b))

  with tf.device('/CPU:0'):
    matmul_sum = tf.add_n(c)

  print(matmul_sum)
 
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op AddN in device /job:localhost/replica:0/task:0/device:CPU:0
tf.Tensor(
[[ 44.  56.]
 [ 98. 128.]], shape=(2, 2), dtype=float32)