使用 GPU

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

TensorFlow 代码和 tf.keras 模型将透明地在单个 GPU 上运行，无需进行任何代码更改。

在多台机器上使用多个 GPU 的最简单方法是使用分布式策略。

本指南适用于尝试过这些方法但发现需要对 TensorFlow 如何使用 GPU 进行细粒度控制的用户。要了解如何调试单 GPU 和多 GPU 场景的性能问题，请参阅优化 TensorFlow GPU 性能指南。

设置

确保已安装最新的 TensorFlow gpu 版本。

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

概述

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:0 和 GPU:0 设备的系统上，会选择 GPU:0 设备来运行 tf.matmul，除非您明确要求在其他设备上运行。

如果 TensorFlow 操作没有相应的 GPU 实现，则该操作将回退到 CPU 设备。例如，由于 tf.cast 只有 CPU 内核，在具有 CPU:0 和 GPU:0 设备的系统上，会选择 CPU:0 设备来运行 tf.cast，即使要求在 GPU:0 设备上运行。

记录设备放置

要找出您的操作和张量分配到的设备，请将 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)

上面的代码将打印一个指示，表明 MatMul 操作已在 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]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)

您将看到现在 a 和 b 被分配到 CPU:0。由于没有为 MatMul 操作显式指定设备，因此 TensorFlow 运行时将根据操作和可用设备（本例中为 GPU:0）选择一个设备，并在需要时自动在设备之间复制张量。

限制 GPU 内存增长

默认情况下，TensorFlow 会将所有 GPU（受 CUDA_VISIBLE_DEVICES 限制）的几乎所有 GPU 内存映射到进程。这样做是为了通过减少内存碎片来更有效地利用设备上相对宝贵的 GPU 内存资源。要将 TensorFlow 限制到特定的一组 GPU，请使用 tf.config.set_visible_devices 方法。

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.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)

在某些情况下，进程可能希望只分配可用内存的一部分，或者只在需要时增长内存使用量。TensorFlow 提供了两种方法来控制这一点。

第一个选项是通过调用 tf.config.experimental.set_memory_growth 来打开内存增长，该方法尝试只分配运行时分配所需的 GPU 内存：它从分配很少的内存开始，随着程序运行和需要更多 GPU 内存，GPU 内存区域将扩展以供 TensorFlow 进程使用。内存不会被释放，因为这会导致内存碎片。要为特定 GPU 打开内存增长，请在分配任何张量或执行任何操作之前使用以下代码。

gpus = tf.config.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.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)

另一种启用此选项的方法是将环境变量 TF_FORCE_GPU_ALLOW_GROWTH 设置为 true。此配置是特定于平台的。

第二种方法是使用 tf.config.set_logical_device_configuration 配置虚拟 GPU 设备，并为 GPU 上分配的总内存设置硬限制。

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.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)

如果您想真正限制 TensorFlow 进程可用的 GPU 内存量，这将很有用。这在本地开发中很常见，因为 GPU 与其他应用程序（如工作站 GUI）共享。

在多 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)

如果您指定的设备不存在，您将收到 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)

使用多个 GPU

为多个 GPU 开发将允许模型随着额外资源的增加而扩展。如果您在具有单个 GPU 的系统上进行开发，可以使用虚拟设备模拟多个 GPU。这使得无需额外资源即可轻松测试多 GPU 设置。

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Create 2 virtual GPUs with 1GB memory each
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024),
         tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.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)

一旦运行时有多个逻辑 GPU 可用，您就可以使用 tf.distribute.Strategy 或手动放置来利用多个 GPU。

使用 `tf.distribute.Strategy`

使用多个 GPU 的最佳实践是使用 tf.distribute.Strategy。以下是一个简单的示例

tf.debugging.set_log_device_placement(True)
gpus = tf.config.list_logical_devices('GPU')
strategy = tf.distribute.MirroredStrategy(gpus)
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))

该程序将在每个 GPU 上运行模型的副本，将输入数据拆分到它们之间，也称为 "数据并行"。

有关分布式策略的更多信息，请查看指南此处.

手动放置

tf.distribute.Strategy 在幕后通过在设备之间复制计算来工作。您可以通过在每个 GPU 上构建模型来手动实现复制。例如

tf.debugging.set_log_device_placement(True)

gpus = tf.config.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)