tensorflow - 设置 tensorflow.keras.mixed_precision.Policy('mixed_float16') 几乎耗尽了所有 GPU 内存

Question

我注意到当我在代码中调用这两行时

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

在我的内核会话的其余部分中，我的 GPU 内存会激增并保持这种状态。在我重新启动内核之前，它不会再次关闭。
当我试图训练一个更大的图像分类模型时，我注意到了这一点，但它内存不足。

这是正常的吗？

我正在使用 tensorflow 版本 2.4.0

标签： tensorflow

Answer 1

这很正常。根据TensorFlow GPU 指南：

默认情况下，TensorFlow 将几乎所有 GPU 的所有 GPU 内存都映射CUDA_VISIBLE_DEVICES到进程可见的所有 GPU 上。这样做是为了通过减少内存碎片更有效地使用设备上相对宝贵的 GPU 内存资源。

使用混合精度允许您将更大的模型加载到 GPU 上。这并不一定意味着 TensorFlow 会保留更少的 GPU 内存。

有两个选项可以限制 GPU 上的内存。这些示例可在此答案开头的链接中找到。

1. 使用tf.config.experimental.set_memory_growth.

gpus = tf.config.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. 用于tf.config.experimental.set_virtual_device_configuration设置 GPU 内存的硬上限。

gpus = tf.config.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)