• tf.config:GPU的使用与分配 *
    • 指定当前程序使用的GPU
    • 设置显存使用策略
    • 单GPU模拟多GPU环境

    tf.config:GPU的使用与分配 *

    指定当前程序使用的GPU

    很多时候的场景是:实验室/公司研究组里有许多学生/研究员需要共同使用一台多GPU的工作站,而默认情况下TensorFlow会使用其所能够使用的所有GPU,这时就需要合理分配显卡资源。

    首先,通过 tf.config.experimental.list_physical_devices ,我们可以获得当前主机上某种特定运算设备类型(如 GPUCPU )的列表,例如,在一台具有4块GPU和一个CPU的工作站上运行以下代码:

    1. gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
    2. cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
    3. print(gpus, cpus)

    输出:

    1. [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'),
    2. PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU'),
    3. PhysicalDevice(name='/physical_device:GPU:2', device_type='GPU'),
    4. PhysicalDevice(name='/physical_device:GPU:3', device_type='GPU')]
    5. [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

    可见,该工作站具有4块GPU:GPU:0GPU:1GPU:2GPU:3 ,以及一个CPU CPU:0

    然后,通过 tf.config.experimental.set_visible_devices ,可以设置当前程序可见的设备范围(当前程序只会使用自己可见的设备,不可见的设备不会被当前程序使用)。例如,如果在上述4卡的机器中我们需要限定当前程序只使用下标为0、1的两块显卡(GPU:0GPU:1),可以使用以下代码:

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

    小技巧

    使用环境变量 CUDA_VISIBLE_DEVICES 也可以控制程序所使用的GPU。假设发现四卡的机器上显卡0,1使用中,显卡2,3空闲,Linux终端输入:

    1. export CUDA_VISIBLE_DEVICES=2,3

    或在代码中加入

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

    即可指定程序只在显卡2,3上运行。

    设置显存使用策略

    默认情况下,TensorFlow将使用几乎所有可用的显存,以避免内存碎片化所带来的性能损失。不过,TensorFlow提供两种显存使用策略,让我们能够更灵活地控制程序的显存使用方式:

    • 仅在需要时申请显存空间(程序初始运行时消耗很少的显存,随着程序的运行而动态申请显存);

    • 限制消耗固定大小的显存(程序不会超出限定的显存大小,若超出的报错)。

    可以通过 tf.config.experimental.set_memory_growth 将GPU的显存使用策略设置为“仅在需要时申请显存空间”。以下代码将所有GPU设置为仅在需要时申请显存空间:

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

    以下代码通过 tf.config.experimental.set_virtual_device_configuration 选项并传入 tf.config.experimental.VirtualDeviceConfiguration 实例,设置TensorFlow固定消耗 GPU:0 的1GB显存(其实可以理解为建立了一个显存大小为1GB的“虚拟GPU”):

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

    提示

    TensorFlow 1.X 的 Graph Execution 下,可以在实例化新的session时传入 tf.compat.v1.ConfigPhoto 类来设置TensorFlow使用显存的策略。具体方式是实例化一个 tf.ConfigProto 类,设置参数,并在创建 tf.compat.v1.Session 时指定Config参数。以下代码通过 allow_growth 选项设置TensorFlow仅在需要时申请显存空间:

    1. config = tf.compat.v1.ConfigProto()
    2. config.gpu_options.allow_growth = True
    3. sess = tf.compat.v1.Session(config=config)

    以下代码通过 per_process_gpu_memory_fraction 选项设置TensorFlow固定消耗40%的GPU显存:

    1. config = tf.compat.v1.ConfigProto()
    2. config.gpu_options.per_process_gpu_memory_fraction = 0.4
    3. tf.compat.v1.Session(config=config)

    单GPU模拟多GPU环境

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

    1. gpus = tf.config.experimental.list_physical_devices('GPU')
    2. tf.config.experimental.set_virtual_device_configuration(
    3. gpus[0],
    4. [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
    5. tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])

    我们在 单机多卡训练 的代码前加入以上代码,即可让原本为多GPU设计的代码在单GPU环境下运行。当输出设备数量时,程序会输出:

    1. Number of devices: 2