Tensorflow如何使用GPU训练（笔记）

Tensorflow和tf.keras 模型可以在单个GPU上透明运行，而无需更改。

注意：（1）需要使用tf.config.experimental.list_physical_devices('GPU')确认使用的tensorflow可以使用GPU。

（2）在一台机器上运行多个GPU，或者在多台机器上运行，最简单的方法是使用分布策略。

确保你的机器已经安装TensorflowGPU版。

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

运行结果：

Num GPUs Available:  2

TensorFlow支持在各种类型的设备上运行计算，包括CPU和GPU。 它们用字符串标识符表示，例如：

（1）"/device:CPU:0": The CPU of your machine.

（2）"/GPU:0": Short-hand notation for the first GPU of your machine that is visible to TensorFlow.

（3）"/job:localhost/replica:0/task:0/device:GPU:1": 你的机器的第二个可用的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)

手动指定设备

如果您希望某个特定操作在您选择的设备上运行，而不是自动为您选择的设备，则可以与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)

限制GPU内存增长

 

默认情况下，TensorFlow会映射该进程可见的几乎所有GPU的所有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来打开内存增长，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内存量，这将很有用。 当GPU与其他应用程序（例如工作站GUI）共享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。

With tf.distribute.Strategy

The best practice for using multiple GPUs is to use tf.distribute.Strategy. Here is a simple example:

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))

 

This program will run a copy of your model on each GPU, splitting the input data between them, also known as "data parallelism".

For more information about distribution strategies, check out the guide here.
 

Manual placement

tf.distribute.Strategy works under the hood by replicating computation across devices. You can manually implement replication by constructing your model on each GPU. For example:

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)