Tensorflow2 GPU训练配置

Tensorflow GPU训练配置

以下教程建立在已经安装好Tensorflow2深度学习环境基础上，进一步管理GPU资源消耗。

查看GPU数量

import tensorflow as tf
# 查看gpu和cpu的数量
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)

可以看到在此我电脑中只有一块GPU。

设置GPU加速

限制使用的GPU，不限制消耗显存的大小

通过tf.config.experimental.set_visible_devices 可以设置当前程序可见的设备范围（当前程序只会使用自己可见的设备，不可见的设备不会被当前程序使用）。

使用部分GPU加速，例如下面使用GPU设备0和1（仅在电脑中有多个GPU设备时使用）

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_visible_devices(devices=gpus[0:2], device_type='GPU')

除了使用以上方法设置可见GPU设备，还可以使用os模块来配置环境变量进行管理。

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"

动态显存申请，仅在需要时申请显存空间

通过tf.config.experimental.set_memory_growth将GPU的显存使用策略设置为“仅在需要时申请显存空间”。

作用：在TensorFlow中，GPU内存默认是一次性分配的，这意味着如果模型占用的内存超过可用内存的限制，将无法运行模型，而会出现OOM（Out Of Memory）错误。为了解决这个问题，TensorFlow提供了函数set_memory_growth，它可以让TensorFlow动态分配GPU内存，只使用所需的GPU内存。总之，使用set_memory_growth函数，可以在程序运行时分配所需的GPU内存，而不是在程序启动时将GPU内存分配给TensorFlow，这样可以避免在运行大型模型时出现内存不足的问题。

以下代码将所有GPU设置为仅在需要时申请显存空间：

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

限制使用的GPU，并且限制使用的显存大小

通过 tf.config.experimental.set_virtual_device_configuration 选项并传入 tf.config.experimental.VirtualDeviceConfiguration 实例，设置TensorFlow固定消耗 GPU:0的1GB显存

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

单GPU模拟多GPU环境

当我们的本地开发环境只有一个GPU，但却需要编写多GPU的程序在工作站上进行训练任务时，TensorFlow为我们提供了一个方便的功能，可以让我们在本地开发环境中建立多个模拟GPU，从而让多GPU的程序调试变得更加方便。以下代码在实体GPU GPU:0 的基础上建立了两个显存均为2GB的虚拟GPU。

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_virtual_device_configuration(
     gpus[0],
     [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
      tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])