 在 TensorFlow.org 上查看 |
 在 Google Colab 中运行 |
 在 GitHub 上查看源代码 |
 下载笔记本 |
无需更改任何代码，TensorFlow 代码以及 tf.keras 模型就可以在单个 GPU 上透明运行。
注：使用 tf.config.experimental.list_physical_devices('GPU') 可以确认 TensorFlow 使用的是 GPU。
在一台或多台机器上，要顺利地在多个 GPU 上运行，最简单的方法是使用分布策略。
本指南适用于已尝试这些方法，但发现需要对 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:0 和 GPU: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 运算在 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)
现在，您会看到已将 a 和 b 分配给 CPU:0。由于没有为 MatMul 运算明确指定设备，TensorFlow 运行时将根据运算和可用的设备选择一个设备（本例中为 GPU:0），并且在需要时会自动在设备之间复制张量。
限制 GPU 内存增长
默认情况下，TensorFlow 会映射进程可见的所有 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 设备，并且设置可在 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 内存量时非常有用。在本地开发中，与其他应用（如工作站 GUI）共享 GPU 时，这是常见做法。
使用多 GPU 系统上的单个 GPU
如果系统上有多个 GPU，则默认情况下会选择具有最小 ID 的 GPU。如果希望在不同的 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 上运行模型的一个副本，并将输入数据拆分到每个 GPU 上，也就是所谓的“数据并行”。
有关分布策略的详细信息，请查阅此处的指南。
手动放置
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)