本笔记参照TensorFlow Guide官方教程,主要是对‘Accelerator-Use a GPU’教程内容翻译和内容结构编排,原文链接:Accelerator-Use a GPU


GPU加速

  • 创建环境(Setup)
  • 一、概览
  • 二、日志设备放置(Logging device placement)
  • 三、手动设备配置(Manual device placement)
  • 四、限制GPU内存增长(Limiting GPU memory growth)
  • 五、在多GPU系统上使用单个GPU
  • 六、使用多GPU
    • 6.1 使用‘tf.distribute.Strategy’
    • 6.2 手动配置


注意: TensorFlow代码和tf.keras模型将透明地运行在一个单独的GPU上,不需要修改代码。

	注意:使用‘tf.config.experimental.list_physical_devices('GPU’)'来确认TensorFlow正在使用GPU。

在一台或多台机器上,在多GPU上运行模型的最简单方法是使用分布策略(Distribution Strategies)。

本指南适用于那些尝试过这些方法并发现需要对TensorFlow如何使用GPU进行微粒度控制的用户。


创建环境(Setup)

确保已经安装最新的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’除非我们显式地要求它运行到另外一个设备上。

二、日志设备放置(Logging device placement)

为了查出我们的操作和张量被配置到哪个设备上,我们可以将‘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)

更多推荐

TensorFlow2.0 Guide官方教程 学习笔记19 -‘Accelerator-Use a GPU‘