![]() |
![]() |
![]() |
![]() |
无需更改任何代码,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)