本笔记参照TensorFlow Guide官方教程,主要是对‘Accelerator-Use a GPU’教程内容翻译和内容结构编排,原文链接:Accelerator-Use a GPU
目录
创建环境(Setup)
一、概览
二、设备配置记录(Logging device placement)
三、手动设备配置(Manual device placement)
四、限制GPU内存增长(Limiting GPU memory growth)
五、在多GPU系统上使用单个GPU
六、使用多GPU
TensorFlow代码和tf.keras模型将透明地运行在一个单独的GPU上,不需要修改代码。
注意:使用‘tf.config.experimental.list_physical_devices('GPU’)'来确认TensorFlow正在使用GPU。
在一台或多台机器上,在多GPU上运行模型的最简单方法是使用分布策略(Distribution Strategies)。
本指南适用于那些尝试过这些方法并发现需要对TensorFlow如何使用GPU进行微粒度控制的用户。
确保已经安装最新的TensorFlow gpu版本
from __future__ import absolute_import, division, print_function, unicode_literals
try:
# %tensorflow_version only exists in Colab.
%tensorflow_version 2.x
except Exception:
pass
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
TensorFlow 2.x selected.
Num GPUs Available: 1
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’上执行。
三、手动设备配置(Manual device placement)
如果希望在自己选择的设备上运行特定的操作,而不是自动为我们选择的操作,我们可以使用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)
tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32)
我们将看到现在a和b被分配给CPU:0。由于MatMul操作没有明确指定设备,TensorFlow运行时将根据操作和可用设备(本例中为GPU:0)选择一个设备,并根据需要在设备之间自动复制张量。
四、限制GPU内存增长(Limiting GPU memory growth)
默认情况下,TensorFlow会将所有GPU(取决于CUDA_VISIBLE_DEVICES)的几乎所有GPU内存映射到进程。这样做是为了通过减少内存碎片更有效地使用设备上相对宝贵的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)
1 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)
Physical devices cannot be modified after being initialized
激活这条选项的另外一种方式是设置环境变量‘TF_FORCE_GPU_ALLOW_GROTH’为‘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)
Virtual devices cannot be modified after being initialized
这是有用的,如果想真正限制的GPU内存量的TensorFlow进程。当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)
/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)
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)
Virtual devices cannot be modified after being initialized
当我们有多个本地GPU用来运行时,我们可以用‘tf.distribute.Strategy’或手动配置这些GPU。
6.1 使用‘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 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
这个程序将在每个GPU上运行我们的模型,把输入数据划分给它们,也叫做‘数据并行’Data_parallelism
更多详情请参考学习笔记18
6.2 手动配置
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
tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32)