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

本笔记参照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进行微粒度控制的用户。

创建环境（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)